{-# LANGUAGE ExtendedDefaultRules #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE UndecidableInstances #-} module TensorSpec (spec) where import Control.Arrow ((&&&)) import Control.Exception.Safe import Numeric.Half import Data.Complex import Data.Int import Data.Word import Test.Hspec import Test.QuickCheck import Torch.DType import Torch.Functional import Torch.Tensor import Torch.TensorFactories import Torch.TensorOptions import Test.QuickCheck.Arbitrary import qualified Data.Vector.Storable as VS import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Generic as VG instance Arbitrary Half where arbitrary = arbitrarySizedFractional shrink = shrinkDecimal instance (Arbitrary a, VS.Storable a) => Arbitrary (VS.Vector a) where arbitrary = do n <- choose (0, 20) -- limit length to at most 20 xs <- vectorOf n arbitrary -- exactly n randomly generated `a`s return (VS.fromList xs) shrink v = [ VS.fromList xs | xs <- shrink (VS.toList v) ] instance (Arbitrary a, VS.Storable a, VG.Vector VU.Vector a) => Arbitrary (VU.Vector a) where arbitrary = do n <- choose (0, 20) -- limit length to at most 20 xs <- vectorOf n arbitrary -- exactly n randomly generated `a`s return (VG.fromList xs) shrink v = [ VG.fromList xs | xs <- shrink (VG.toList v) ] spec :: Spec spec = do describe "TensorLike" $ do it "TensorLike Bool" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Bool) it "TensorLike Word8" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Word8) it "TensorLike Int8" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Int8) it "TensorLike Int16" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Int16) it "TensorLike Int32" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Int32) it "TensorLike Int" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Int) it "TensorLike Int64" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Int64) it "TensorLike Half" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Half) it "TensorLike Float" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Float) it "TensorLike Double" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Double) it "TensorLike ComplexHalf" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Complex Half) it "TensorLike Complex Float" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Complex Float) it "TensorLike Complex Double" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: Complex Double) it "TensorLike Storable Vector Float" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: VS.Vector Float) it "TensorLike Storable Vector Double" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: VS.Vector Double) it "TensorLike Unboxed Vector Double" $ property $ \x -> asValue (asTensor x) `shouldBe` (x :: VU.Vector Double) it "Compare internal expression of c++ with Storable expression of haskell" $ do show (asTensor [True, False, True, False]) `shouldBe` "Tensor Bool [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Word8])) `shouldBe` "Tensor UInt8 [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Int8])) `shouldBe` "Tensor Int8 [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Int16])) `shouldBe` "Tensor Int16 [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Int32])) `shouldBe` "Tensor Int32 [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Int])) `shouldBe` "Tensor Int64 [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Int64])) `shouldBe` "Tensor Int64 [4] [ 1, 0, 1, 0]" show (asTensor ([1, 0, 1, 0] :: [Float])) `shouldBe` "Tensor Float [4] [ 1.0000 , 0.0000, 1.0000 , 0.0000]" show (asTensor ([1, 0, 1, 0] :: [Double])) `shouldBe` "Tensor Double [4] [ 1.0000 , 0.0000, 1.0000 , 0.0000]" show (asTensor ([[]] :: [[Int]])) `shouldBe` "Tensor Int64 [1,0] [[]]" show (asTensor ([[1]] :: [[Int]])) `shouldBe` "Tensor Int64 [1,1] [[ 1]]" it "TensorLike [Bool]" $ property $ \(NonEmpty (x :: [Bool])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` (if (head x) then 1 else 0) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Word8]" $ property $ \(NonEmpty (x :: [Word8])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Int8]" $ property $ \(NonEmpty (x :: [Int8])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Int16]" $ property $ \(NonEmpty (x :: [Int16])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Int32]" $ property $ \(NonEmpty (x :: [Int32])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Int]" $ property $ \(NonEmpty (x :: [Int])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Int64]" $ property $ \(NonEmpty (x :: [Int64])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Half]" $ property $ \(NonEmpty (x :: [Half])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` realToFrac (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Float]" $ property $ \(NonEmpty (x :: [Float])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` realToFrac (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Double]" $ property $ \(NonEmpty (x :: [Double])) -> do asValue (asTensor x) `shouldBe` x toDouble (select 0 0 (asTensor x)) `shouldBe` realToFrac (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Complex Half]" $ property $ \(NonEmpty (x :: [Complex Half])) -> do asValue (asTensor x) `shouldBe` x asValue (select 0 0 (asTensor x)) `shouldBe` (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Complex Float]" $ property $ \(NonEmpty (x :: [Complex Float])) -> do asValue (asTensor x) `shouldBe` x asValue (select 0 0 (asTensor x)) `shouldBe` (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "TensorLike [Complex Double]" $ property $ \(NonEmpty (x :: [Complex Double])) -> do asValue (asTensor x) `shouldBe` x asValue (select 0 0 (asTensor x)) `shouldBe` (head x) shape (asTensor x) `shouldBe` [length x] let xx = replicate 5 x asValue (asTensor xx) `shouldBe` xx let xxx = replicate 3 xx asValue (asTensor xxx) `shouldBe` xxx it "invalid cast of TensorLike a" $ do let x = asTensor (10 :: Int) (dtype x) `shouldBe` Int64 (print (asValue x :: Double)) `shouldThrow` anyException it "invalid cast of TensorLike [a]" $ do let x = asTensor ([0 .. 10] :: [Int]) (print (asValue x :: [Double])) `shouldThrow` anyException it "lists having different length" $ do (print (asTensor ([[1], [1, 2]] :: [[Double]]))) `shouldThrow` anyException it "cast of Tensor" $ do let x = asTensor ([0 .. 10] :: [Int]) (dtype (toType Float x)) `shouldBe` Float describe "indexing" $ do it "pick up a value" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = x ! (1, 0, 2) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([], 8 :: Int)) it "pick up a bottom tensor" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = x ! (1, 0) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([3], [6 :: Int, 7, 8])) it "make a slice of bottom values" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = x ! ((), (), 1) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2], [[1 :: Int, 4], [7, 10]])) it "ellipsis" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = x ! (Ellipsis, 1) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2], [[1 :: Int, 4], [7, 10]])) it "make a slice via muliple slices" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = x ! ((), (Slice (1, None))) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 1, 3], [[[3 :: Int, 4, 5]], [[9, 10, 11]]])) it "make a slice via muliple slices" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = x ! ((), (Slice (1, None)), (Slice (0, 1))) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 1, 1], [[[3 :: Int]], [[9]]])) describe "masked fill" $ do it "Fill a value" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = maskedFill x (1, 0, 2) (9 :: Int) (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2, 3], [[[0, 1, 2], [3, 4, 5]], [[6, 7, 9], [9, 10, 11]]] :: [[[Int]]])) it "Fill a bottom tensor" $ do let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = maskedFill x (1, 0) [8 :: Int, 8, 8] (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2, 3], [[[0, 1, 2], [3, 4, 5]], [[8, 8, 8], [9, 10, 11]]] :: [[[Int]]])) it "masked fill by boolean" $ do let m = asTensor ([[[True, True, False], [False, False, False]], [[True, False, False], [False, False, False]]] :: [[[Bool]]]) x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]]) r = maskedFill x m [8 :: Int, 8, 8] (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2, 3], [[[8, 8, 2], [3, 4, 5]], [[8, 7, 8], [9, 10, 11]]] :: [[[Int]]]))