hasktorch 0.2.1.6 → 0.2.1.7

raw patch · 16 files changed

+397/−9 lines, 16 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• hasktorch.cabal +5/−1
• src/Torch/Functional.hs +79/−0
• src/Torch/Functional/Internal.hs +4/−2
• src/Torch/Tensor.hs +34/−0
• src/Torch/Typed/Auxiliary.hs +10/−2
• src/Torch/Typed/Functional.hs +31/−2
• src/Torch/Typed/Tensor.hs +32/−0
• test/FunctionalSpec.hs +13/−0
• test/Torch/Typed/AuxiliarySpec.hs +25/−1
• test/Torch/Typed/FactoriesSpec.hs +10/−0
• test/Torch/Typed/FunctionalSpec0.hs +41/−1
• test/Torch/Typed/FunctionalSpec1.hs +21/−0
• test/Torch/Typed/FunctionalSpec2.hs +36/−0
• test/Torch/Typed/OptimSpec.hs +21/−0
• test/Torch/Typed/TensorSpec0.hs +19/−0
• test/Torch/Typed/TensorSpec1.hs +16/−0

@@ -1,6 +1,6 @@ cabal-version:       3.0 name:                hasktorch-version:             0.2.1.6+version:             0.2.1.7 synopsis:            Haskell bindings to libtorch, supporting both typed and untyped tensors. description:         Hasktorch is a library for tensors and neural networks in Haskell. It is an independent open source community project which leverages the core C++ libraries shared by PyTorch. homepage:            https://github.com/hasktorch/hasktorch#readme@@ -103,6 +103,8 @@  hs-source-dirs:      src  default-language:    Haskell2010  ghc-options:         -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver -fplugin GHC.TypeLits.Extra.Solver -fconstraint-solver-iterations=0 -fplugin GHC.NotExport.Plugin+ if os(darwin)+  cpp-options:        -D__APPLE__  build-depends:       async >= 2.2.5 && < 2.3                     , base >= 4.7 && < 5                     , libtorch-ffi == 2.0.1.*@@ -190,6 +192,8 @@                     , IndexSpec   default-language: Haskell2010   ghc-options:         -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver -fplugin GHC.TypeLits.Extra.Solver -fconstraint-solver-iterations=0+  if os(darwin)+   cpp-options:       -D__APPLE__   build-depends:      base                     , ghc-typelits-extra                     , ghc-typelits-knownnat

@@ -21,6 +21,7 @@     Internal.logdet,     Internal.lstsq,     Internal.mv,+    Internal.scaled_dot_product_attention,     Internal.sumWithDimnames,   ) where@@ -728,6 +729,84 @@ eq a b = unsafePerformIO $ cast2 ATen.eq_tt a b  (==.) = eq++-- | Computes input > scalar element-wise.+-- The second argument is a scalar value that is compared against each element of the tensor.+gtScalar ::+  -- | input+  Tensor ->+  -- | scalar+  Float ->+  -- | output+  Tensor+gtScalar = Internal.gtScalar++(.>) = gtScalar++-- | Computes input < scalar element-wise.+-- The second argument is a scalar value that is compared against each element of the tensor.+ltScalar ::+  -- | input+  Tensor ->+  -- | scalar+  Float ->+  -- | output+  Tensor+ltScalar = Internal.ltScalar++(.<) = ltScalar++-- | Computes input >= scalar element-wise.+-- The second argument is a scalar value that is compared against each element of the tensor.+geScalar ::+  -- | input+  Tensor ->+  -- | scalar+  Float ->+  -- | output+  Tensor+geScalar = Internal.geScalar++(.>=) = geScalar++-- | Computes input <= scalar element-wise.+-- The second argument is a scalar value that is compared against each element of the tensor.+leScalar ::+  -- | input+  Tensor ->+  -- | scalar+  Float ->+  -- | output+  Tensor+leScalar = Internal.leScalar++(.<=) = leScalar++-- | Computes input == scalar element-wise.+-- The second argument is a scalar value that is compared against each element of the tensor.+eqScalar ::+  -- | input+  Tensor ->+  -- | scalar+  Float ->+  -- | output+  Tensor+eqScalar = Internal.eqScalar++(.==) = eqScalar++-- | Computes input /= scalar element-wise.+-- The second argument is a scalar value that is compared against each element of the tensor.+neScalar ::+  -- | input+  Tensor ->+  -- | scalar+  Float ->+  -- | output+  Tensor+neScalar = Internal.neScalar++(./=) = neScalar  -- | Returns a new tensor with the elements of input at the given indices. The input tensor is treated as if it were viewed as a 1-D tensor. The result takes the same shape as the indices. take ::

@@ -6440,11 +6440,13 @@   :: Tensor -- ^ query   -> Tensor -- ^ key   -> Tensor -- ^ value-  -> Tensor -- ^ attn_mask+  -> Maybe Tensor -- ^ attn_mask   -> Double -- ^ dropout_p   -> Bool -- ^ is_causal+  -> Double -- ^ scale+  -> Bool -- ^ enable_gqa   -> Tensor-scaled_dot_product_attention _query _key _value _attn_mask _dropout_p _is_causal = unsafePerformIO $ (cast6 ATen.scaled_dot_product_attention_ttttdb) _query _key _value _attn_mask _dropout_p _is_causal+scaled_dot_product_attention _query _key _value _attn_mask _dropout_p _is_causal _scale _enable_gqa = unsafePerformIO $ (cast8 ATen.scaled_dot_product_attention_tttqdbdb) _query _key _value _attn_mask _dropout_p _is_causal _scale _enable_gqa  special_airy_ai   :: Tensor -- ^ x

@@ -49,6 +49,7 @@ import qualified Torch.Internal.Managed.Type.StdArray as ATen import qualified Torch.Internal.Managed.Type.StdString as ATen import qualified Torch.Internal.Managed.Type.Tensor as ATen+import qualified Torch.Internal.Managed.Type.StdOptional as ATen import qualified Torch.Internal.Managed.Type.TensorIndex as ATen import qualified Torch.Internal.Managed.Type.TensorOptions as ATen import qualified Torch.Internal.Managed.Type.Extra as ATen@@ -77,6 +78,39 @@  toImmutable :: MutableTensor -> IO Tensor toImmutable (MutableTensor tensor) = cast1 ATen.detach_t tensor++type ATenOptionalTensor = ForeignPtr (ATen.StdOptional ATen.Tensor)++-- do not use the constructor+newtype OptionalTensor = UnsafeOptional ATenOptionalTensor++instance Castable OptionalTensor ATenOptionalTensor where+  cast (UnsafeOptional aten_optional_tensor) f = f aten_optional_tensor+  uncast aten_optional_tensor f = f $ UnsafeOptional aten_optional_tensor++instance Castable (Maybe Tensor) OptionalTensor where+  cast Nothing f = do+    ptr <- ATen.stdOptionalTensor_empty+    f (UnsafeOptional ptr)+  cast (Just tensor) f = do+    cast tensor $ \atenTensor -> do+      ptr <- ATen.stdOptionalTensor_create atenTensor+      f (UnsafeOptional ptr)+  uncast (UnsafeOptional ptr) f = do+    hasValue <- ATen.stdOptionalTensor_has_value ptr+    if hasValue /= 0+      then do+        atenTensor <- ATen.stdOptionalTensor_value ptr+        uncast atenTensor $ \tensor -> f (Just tensor)+      else f Nothing++instance Castable (Maybe Tensor) ATenOptionalTensor where+  cast maybeTensor f = do+    cast maybeTensor $ \(optTensor :: OptionalTensor) -> do+      cast optTensor f+  uncast ptr f = do+    uncast ptr $ \(optTensor :: OptionalTensor) -> do+      uncast optTensor f  -------------------------------------------------------------------------------- -- Basic tensor properties

@@ -328,6 +328,10 @@   DTypeIsNotHalf '(deviceType, _) D.Half = UnsupportedDTypeForDevice deviceType D.Half   DTypeIsNotHalf _ _ = () +type family DTypeIsNotDouble (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where+  DTypeIsNotDouble '(deviceType, _) D.Double = UnsupportedDTypeForDevice deviceType D.Double+  DTypeIsNotDouble _ _ = ()+ type family DTypeIsNotBool (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where   DTypeIsNotBool '(deviceType, _) D.Bool = UnsupportedDTypeForDevice deviceType D.Bool   DTypeIsNotBool _ _ = ()@@ -347,11 +351,11 @@     ( DTypeIsFloatingPoint '( 'D.CPU, 0) dtype,       DTypeIsNotHalf '( 'D.CPU, 0) dtype     )+  StandardFloatingPointDTypeValidation '( 'D.CUDA, deviceIndex) dtype = DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype   StandardFloatingPointDTypeValidation '( 'D.MPS, 0) dtype =     ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,-      DTypeIsNotHalf '( 'D.MPS, 0) dtype+      DTypeIsNotDouble '( 'D.MPS, 0) dtype     )-  StandardFloatingPointDTypeValidation '( 'D.CUDA, deviceIndex) dtype = DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype   StandardFloatingPointDTypeValidation '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype  type family StandardDTypeValidation (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where@@ -360,6 +364,10 @@       DTypeIsNotHalf '( 'D.CPU, 0) dtype     )   StandardDTypeValidation '( 'D.CUDA, deviceIndex) dtype = DTypeIsNotBool '( 'D.CUDA, deviceIndex) dtype+  StandardDTypeValidation '( 'D.MPS, 0) dtype =+    ( DTypeIsNotBool '( 'D.MPS, 0) dtype,+      DTypeIsNotDouble '( 'D.MPS, 0) dtype+    )   StandardDTypeValidation '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype  

@@ -572,6 +572,23 @@   Tensor device dtype shape divScalar a input = unsafePerformIO $ ATen.cast2 ATen.Managed.div_ts input a +-- | divScalar'+-- TODO: what dtypes is this defined for?+-- TODO: what scalar types is this defined for?+--+-- >>> dtype &&& shape $ divScalar 2 (ones :: CPUTensor 'D.Float '[2,2])+-- (Float,[2,2])+divScalar' ::+  forall a shape dtype device.+  D.Scalar a =>+  -- | tensor input+  Tensor device dtype shape ->+  -- | scalar input+  a ->+  -- | output+  Tensor device dtype shape+divScalar' input a = a `mulScalar` reciprocal input+ -- | powScalar -- TODO: probably only defined for floating point tensors, or maybe numeric type is lifted? --@@ -4547,8 +4564,20 @@   Tensor device dtype shape' squeezeDim input = unsafePerformIO $ ATen.cast2 ATen.Managed.squeeze_tl input (natValI @dim) --- where' :: Tensor device dtype shape -> Tensor device dtype shape -> Tensor device dtype shape -> Tensor device dtype shape--- where' _condition _input _other = unsafePerformIO $ (ATen.cast3 ATen.Managed.where_ttt) _condition _input _other+-- | return a tensor of elements selected from either input or other, depending on condition.+where' ::+  forall shape shape' shape'' shape''' shape'''' dtype device.+  ( shape'' ~ Broadcast shape shape',+    shape'''' ~ Broadcast shape''' shape''+  ) =>+  -- | condition+  Tensor device 'D.Bool shape ->+  -- | input+  Tensor device dtype shape' ->+  -- | other+  Tensor device dtype shape''' ->+  Tensor device dtype shape''''+where' _condition _input _other = unsafePerformIO $ (ATen.cast3 ATen.Managed.where_ttt) _condition _input _other  -- where_ :: Tensor device dtype shape -> [Tensor device dtype shape] -- where_ _condition = unsafePerformIO $ (ATen.cast1 ATen.Managed.where_t) _condition

@@ -451,6 +451,38 @@ (==.) = eq (/=.) = ne +gtScalar,+  ltScalar,+  geScalar,+  leScalar,+  eqScalar,+  neScalar,+  (.>),+  (.<),+  (.>=),+  (.<=),+  (.==),+  (./=) ::+    forall shape dtype device.+    ( ComparisonDTypeIsValid device dtype,+      StandardFloatingPointDTypeValidation device dtype+    ) =>+    Tensor device dtype shape ->+    Float ->+    Tensor device 'D.Bool shape+gtScalar a s = UnsafeMkTensor $ D.gtScalar (toDynamic a) s+ltScalar a s = UnsafeMkTensor $ D.ltScalar (toDynamic a) s+geScalar a s = UnsafeMkTensor $ D.geScalar (toDynamic a) s+leScalar a s = UnsafeMkTensor $ D.leScalar (toDynamic a) s+eqScalar a s = UnsafeMkTensor $ D.eqScalar (toDynamic a) s+neScalar a s = UnsafeMkTensor $ D.neScalar (toDynamic a) s+(.>) = gtScalar+(.<) = ltScalar+(.>=) = geScalar+(.<=) = leScalar+(.==) = eqScalar+(./=) = neScalar+ type family ComputeMatMul (reversedShape :: [Nat]) (reversedShape' :: [Nat]) :: Maybe [Nat] where   ComputeMatMul (k ': '[]) (k ': '[]) = Just '[]   ComputeMatMul (k ': '[]) (m ': k ': reversedBroadcastShape') = AppendToMaybe m (ComputeBroadcast '[] reversedBroadcastShape')

@@ -300,3 +300,16 @@       (asTensor ([[[0.1, 0.2, 0.7]]] :: [[[Float]]]))       (asTensor ([2] :: [Int]))       `shouldBe` asTensor (-0.7 :: Float)+  it "scaled_dot_product_attention with Nothing mask" $ do+    let query = ones' [2, 4, 8]+        key = ones' [2, 4, 8]+        value = ones' [2, 4, 8]+        result = scaled_dot_product_attention query key value Nothing 0.0 True 1.0 False+    shape result `shouldBe` [2, 4, 8]+  it "scaled_dot_product_attention with Just mask" $ do+    let query = ones' [2, 4, 8]+        key = ones' [2, 4, 8]+        value = ones' [2, 4, 8]+        attn_mask = ones' [2, 4, 4]+        result = scaled_dot_product_attention query key value (Just attn_mask) 0.0 False 1.0 False+    shape result `shouldBe` [2, 4, 8]

@@ -9,11 +9,12 @@  import Data.Proxy import System.IO.Unsafe+import System.Environment (lookupEnv) import Test.Hspec (Spec, shouldBe) import Test.QuickCheck () import qualified Torch as Torch (device, dtype, shape) import Torch.Internal.Cast (cast0)-import Torch.Internal.Managed.Type.Context (hasCUDA)+import Torch.Internal.Managed.Type.Context (hasCUDA, hasMPS) import Torch.Typed  instance Semigroup Spec where@@ -37,6 +38,9 @@ allFloatingPointDTypes :: _ allFloatingPointDTypes = withHalf standardFloatingPointDTypes +mpsFloatingPointDTypes :: _+mpsFloatingPointDTypes = Proxy @'Float :. HNil+ standardFloatingPointDTypes :: _ standardFloatingPointDTypes = Proxy @'Float :. Proxy @'Double :. HNil @@ -61,18 +65,38 @@     :. Proxy @'Int64     :. standardFloatingPointDTypes +mpsDTypes :: _+mpsDTypes =+  Proxy @'UInt8+    :. Proxy @'Int8+    :. Proxy @'Int16+    :. Proxy @'Int32+    :. Proxy @'Int64+    :. mpsFloatingPointDTypes+ cpu :: _ cpu = Proxy @'( 'CPU, 0)  cuda0 :: _ cuda0 = Proxy @'( 'CUDA, 0) +mps :: _+mps = Proxy @'( 'MPS, 0)+ availableDevices :: [Device] availableDevices =   let hasCuda = unsafePerformIO $ cast0 hasCUDA+      hasMps = unsafePerformIO $ cast0 hasMPS+      hasMpsFallback = unsafePerformIO $ do+        env <- lookupEnv "PYTORCH_ENABLE_MPS_FALLBACK"+        return $ env == Just "1"    in [Device {deviceType = CPU, deviceIndex = 0}]         <> ( if hasCuda                then [Device {deviceType = CUDA, deviceIndex = 0}]+               else mempty+           )+        <> ( if hasMps && hasMpsFallback+               then [Device {deviceType = MPS, deviceIndex = 0}]                else mempty            ) 

@@ -68,6 +68,8 @@                 hfoldrM @IO simpleFactoriesSpec () (hattach cpu (hproduct allDTypes standardShapes))               Device {deviceType = CUDA, deviceIndex = 0} ->                 hfoldrM @IO simpleFactoriesSpec () (hattach cuda0 (hproduct allDTypes standardShapes))+              Device {deviceType = MPS, deviceIndex = 0} ->+                hfoldrM @IO simpleFactoriesSpec () (hattach mps (hproduct mpsDTypes standardShapes))       it "ones" $ dispatch ZerosSpec       it "zeros" $ dispatch OnesSpec       it "full" $ dispatch FullSpec@@ -78,6 +80,8 @@                 hfoldrM @IO randomFactoriesSpec () (hattach cpu (hproduct standardFloatingPointDTypes standardShapes))               Device {deviceType = CUDA, deviceIndex = 0} ->                 hfoldrM @IO randomFactoriesSpec () (hattach cuda0 (hproduct allFloatingPointDTypes standardShapes))+              Device {deviceType = MPS, deviceIndex = 0} ->+                hfoldrM @IO randomFactoriesSpec () (hattach mps (hproduct mpsFloatingPointDTypes standardShapes))       it "rand" $ dispatch RandSpec       it "randn" $ dispatch RandnSpec     describe "advanced factories" $ do@@ -88,10 +92,16 @@         Device {deviceType = CUDA, deviceIndex = 0} -> do           let t = linspace @3 @'( 'CUDA, 0) (1 :: Int) (3 :: Int)           checkDynamicTensorAttributes t+        Device {deviceType = MPS, deviceIndex = 0} -> do+          let t = linspace @3 @'( 'MPS, 0) (1 :: Int) (3 :: Int)+          checkDynamicTensorAttributes t       it "eyeSquare" $ case device of         Device {deviceType = CPU, deviceIndex = 0} -> do           let t = eyeSquare @10 @'Float @'( 'CPU, 0)           checkDynamicTensorAttributes t         Device {deviceType = CUDA, deviceIndex = 0} -> do           let t = eyeSquare @10 @'Float @'( 'CUDA, 0)+          checkDynamicTensorAttributes t+        Device {deviceType = MPS, deviceIndex = 0} -> do+          let t = eyeSquare @10 @'Float @'( 'MPS, 0)           checkDynamicTensorAttributes t

@@ -537,12 +537,16 @@               hfoldrM @IO unaryAllDTypesSpec () (hattach cpu (hproduct allDTypes standardShapes))             Device {deviceType = CUDA, deviceIndex = 0} ->               hfoldrM @IO unaryAllDTypesSpec () (hattach cuda0 (hproduct allDTypes standardShapes))+            Device {deviceType = MPS, deviceIndex = 0} ->+              hfoldrM @IO unaryAllDTypesSpec () (hattach mps (hproduct mpsDTypes standardShapes))        it "abs" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO AbsSpec () (hattach cpu (hproduct standardDTypes standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO AbsSpec () (hattach cuda0 (hproduct standardDTypes standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO AbsSpec () (hattach mps (hproduct mpsDTypes standardShapes))       it "sign" $ dispatch SignSpec       it "onesLike" $ dispatch OnesLikeSpec       it "zerosLike" $ dispatch ZerosLikeSpec@@ -554,6 +558,8 @@               hfoldrM @IO unaryStandardFloatingPointDTypesSpec () (hattach cpu (hproduct standardFloatingPointDTypes standardShapes))             Device {deviceType = CUDA, deviceIndex = 0} ->               hfoldrM @IO unaryStandardFloatingPointDTypesSpec () (hattach cuda0 (hproduct allFloatingPointDTypes standardShapes))+            Device {deviceType = MPS, deviceIndex = 0} ->+              hfoldrM @IO unaryStandardFloatingPointDTypesSpec () (hattach mps (hproduct mpsFloatingPointDTypes standardShapes))        it "frac" $ dispatch FracSpec       it "ceil" $ dispatch CeilSpec@@ -579,11 +585,15 @@           hfoldrM @IO MishSpec () (hattach cpu (hproduct standardFloatingPointDTypes standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO MishSpec () (hattach cuda0 (hproduct allFloatingPointDTypes standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO MishSpec () (hattach mps (hproduct mpsFloatingPointDTypes standardShapes))       it "gelu" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO GeluSpec () (hattach cpu (hproduct standardFloatingPointDTypes standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO GeluSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO GeluSpec () (hattach mps (hproduct mpsFloatingPointDTypes standardShapes))       it "leakyRelu" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO@@ -595,6 +605,11 @@             LeakyReluSpec             ()             (hattach cuda0 (hproduct scalarParams (hproduct standardFloatingPointDTypes standardShapes)))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO+            LeakyReluSpec+            ()+            (hattach mps (hproduct scalarParams (hproduct mpsFloatingPointDTypes standardShapes)))       it "elu" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO@@ -606,6 +621,11 @@             ELUSpec             ()             (hattach cuda0 (hproduct (hzip scalarParams (hzip scalarParams scalarParams)) (hproduct standardFloatingPointDTypes standardShapes)))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO+            ELUSpec+            ()+            (hattach mps (hproduct (hzip scalarParams (hzip scalarParams scalarParams)) (hproduct mpsFloatingPointDTypes standardShapes)))       it "sigmoid" $ dispatch SigmoidSpec       it "logSigmoid" $ dispatch LogSigmoidSpec @@ -629,6 +649,8 @@           hfoldrM @IO ToDTypeSpec () (hattach cpu (hproduct (hproduct allDTypes allDTypes) standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO ToDTypeSpec () (hattach cuda0 (hproduct (hproduct allDTypes allDTypes) standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO ToDTypeSpec () (hattach mps (hproduct (hproduct mpsDTypes mpsDTypes) standardShapes))      describe "aggregation" $ do       it "sumAll" $ case device of@@ -636,6 +658,8 @@           hfoldrM @IO SumAllSpec () (hattach cpu (hproduct almostAllDTypes standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO SumAllSpec () (hattach cuda0 (hproduct allDTypes standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO SumAllSpec () (hattach mps (hproduct mpsDTypes standardShapes))       it "sumDim" $ do         let sumDimDims = Proxy @0 :. Proxy @1 :. HNil             sumDimShapes = Proxy @'[1, 0] :. Proxy @'[2, 3] :. HNil@@ -644,6 +668,8 @@             hfoldrM @IO SumDimSpec () (hproduct sumDimDims (hattach cpu (hproduct almostAllDTypes sumDimShapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO SumDimSpec () (hproduct sumDimDims (hattach cuda0 (hproduct allDTypes sumDimShapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO SumDimSpec () (hproduct sumDimDims (hattach mps (hproduct mpsDTypes sumDimShapes)))       do         let shapes = (Proxy :: Proxy ('[] :: [Nat])) :. Proxy @'[1] :. Proxy @'[2, 3] :. HNil             dispatch spec = case device of@@ -651,6 +677,8 @@                 hfoldrM @IO spec () (hattach cpu (hproduct almostAllDTypes shapes))               Device {deviceType = CUDA, deviceIndex = 0} ->                 hfoldrM @IO spec () (hattach cuda0 (hproduct allDTypes shapes))+              Device {deviceType = MPS, deviceIndex = 0} ->+                hfoldrM @IO spec () (hattach mps (hproduct mpsDTypes shapes))         it "min" $ dispatch MinSpec         it "max" $ dispatch MaxSpec       it "meanAll" $ do@@ -660,6 +688,8 @@             hfoldrM @IO MeanAllSpec () (hattach cpu (hproduct standardFloatingPointDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MeanAllSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MeanAllSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "meanDim" $ do         let dims = Proxy @0 :. Proxy @1 :. HNil             shapes = Proxy @'[1, 3] :. Proxy @'[2, 3] :. HNil@@ -668,6 +698,8 @@             hfoldrM @IO MeanDimSpec () (hproduct dims (hattach cpu (hproduct standardFloatingPointDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MeanDimSpec () (hproduct dims (hattach cuda0 (hproduct standardFloatingPointDTypes shapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MeanDimSpec () (hproduct dims (hattach mps (hproduct mpsFloatingPointDTypes shapes)))       it "mean" $ do         let dims = Proxy @0 :. Proxy @1 :. HNil             keepOrDropDims = Proxy @KeepDim :. Proxy @DropDim :. HNil@@ -677,6 +709,8 @@             hfoldrM @IO MeanSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cpu (hproduct standardFloatingPointDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MeanSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cpu (hproduct standardFloatingPointDTypes shapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MeanSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cpu (hproduct mpsFloatingPointDTypes shapes)))       it "medianAll" $ do         let shapes = (Proxy :: Proxy ('[] :: [Nat])) :. Proxy @'[1] :. Proxy @'[2, 3] :. HNil         case device of@@ -684,6 +718,8 @@             hfoldrM @IO MedianAllSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MedianAllSpec () (hattach cuda0 (hproduct (withHalf standardDTypes) shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MedianAllSpec () (hattach mps (hproduct (withHalf mpsDTypes) shapes))       it "medianDim" $ do         let dims = Proxy @0 :. Proxy @1 :. HNil             shapes = Proxy @'[1, 17, 1] :. Proxy @'[2, 3] :. HNil@@ -692,6 +728,8 @@             hfoldrM @IO MedianDimSpec () (hproduct dims (hattach cpu (hproduct standardDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MedianDimSpec () (hproduct dims (hattach cuda0 (hproduct (withHalf standardDTypes) shapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MedianDimSpec () (hproduct dims (hattach mps (hproduct (withHalf mpsDTypes) shapes)))       it "median" $ do         let dims = Proxy @0 :. Proxy @1 :. HNil             keepOrDropDims = Proxy @KeepDim :. Proxy @DropDim :. HNil@@ -701,6 +739,8 @@             hfoldrM @IO MedianSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cpu (hproduct standardDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MedianSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cuda0 (hproduct (withHalf standardDTypes) shapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MedianSpec () (hproduct (hproduct dims keepOrDropDims) (hattach mps (hproduct (withHalf mpsDTypes) shapes)))       it "mode" $ do         let dims = Proxy @0 :. Proxy @1 :. HNil             keepOrDropDims = Proxy @KeepDim :. Proxy @DropDim :. HNil@@ -710,4 +750,4 @@             hfoldrM @IO ModeSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cpu (hproduct standardDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO ModeSpec () (hproduct (hproduct dims keepOrDropDims) (hattach cuda0 (hproduct (withHalf standardDTypes) shapes)))-+          Device {deviceType = MPS, deviceIndex = 0} -> pure ()

@@ -222,11 +222,15 @@             hfoldrM @IO NarrowSpec () (hproduct dims (hproduct narrowStarts (hproduct narrowLengths (hattach cpu (hproduct standardFloatingPointDTypes narrowShapes)))))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO NarrowSpec () (hproduct dims (hproduct narrowStarts (hproduct narrowLengths (hattach cuda0 (hproduct allFloatingPointDTypes narrowShapes)))))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO NarrowSpec () (hproduct dims (hproduct narrowStarts (hproduct narrowLengths (hattach mps (hproduct mpsFloatingPointDTypes narrowShapes)))))       it "squeezeAll" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO SqueezeAllSpec () (hattach cpu (hproduct allDTypes standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO SqueezeAllSpec () (hattach cuda0 (hproduct allDTypes standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO SqueezeAllSpec () (hattach mps (hproduct mpsDTypes standardShapes))       it "transpose" $ do         let dims =               hzip@@ -238,11 +242,15 @@             hfoldrM @IO TransposeSpec () (hproduct dims (hattach cpu (hproduct allDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO TransposeSpec () (hproduct dims (hattach cuda0 (hproduct allDTypes shapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO TransposeSpec () (hproduct dims (hattach mps (hproduct mpsDTypes shapes)))       it "transpose2d" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO Transpose2DSpec () (hattach cpu (hproduct allDTypes (Proxy @'[2, 3] :. HNil)))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO Transpose2DSpec () (hattach cuda0 (hproduct allDTypes (Proxy @'[2, 3] :. HNil)))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO Transpose2DSpec () (hattach mps (hproduct mpsDTypes (Proxy @'[2, 3] :. HNil)))       it "diag" $ do         let vectorShapes = Proxy @'[0] :. Proxy @'[1] :. Proxy @'[2] :. HNil             emptyShapes = Proxy @'[0, 0] :. Proxy @'[0, 1] :. Proxy @'[1, 0] :. HNil@@ -258,6 +266,10 @@             hfoldrM @IO DiagSpec () (hproduct (hproduct tris indexes) (hattach cuda0 (hproduct (withHalf standardDTypes) standardShapes)))             hfoldrM @IO DiagSpec () (hproduct (hproduct tris indexes) (hattach cuda0 (hproduct (withHalf standardDTypes) vectorShapes)))             hfoldrM @IO DiagSpec () (hproduct (hproduct tris indexes') (hattach cuda0 (hproduct (withHalf standardDTypes) emptyShapes)))+          Device {deviceType = MPS, deviceIndex = 0} -> do+            hfoldrM @IO DiagSpec () (hproduct (hproduct tris indexes) (hattach mps (hproduct (withHalf mpsDTypes) standardShapes)))+            hfoldrM @IO DiagSpec () (hproduct (hproduct tris indexes) (hattach mps (hproduct (withHalf mpsDTypes) vectorShapes)))+            hfoldrM @IO DiagSpec () (hproduct (hproduct tris indexes') (hattach mps (hproduct (withHalf mpsDTypes) emptyShapes)))       it "diagEmbed" $ do         let shapes =               standardShapes@@ -279,6 +291,9 @@           Device {deviceType = CUDA, deviceIndex = 0} -> do             hfoldrM @IO DiagEmbedSpec () (hproduct (hproduct indexes dims) (hattach cuda0 (hproduct (withHalf standardDTypes) shapes)))             hfoldrM @IO DiagEmbedSpec () (hproduct (hproduct indexes allDims) (hattach cuda0 (hproduct (withHalf standardDTypes) standardShapes)))+          Device {deviceType = MPS, deviceIndex = 0} -> do+            hfoldrM @IO DiagEmbedSpec () (hproduct (hproduct indexes dims) (hattach mps (hproduct (withHalf mpsDTypes) shapes)))+            hfoldrM @IO DiagEmbedSpec () (hproduct (hproduct indexes allDims) (hattach mps (hproduct (withHalf mpsDTypes) standardShapes)))       it "diagflat" $ do         let shapes =               standardShapes@@ -296,6 +311,8 @@             hfoldrM @IO DiagflatSpec () (hproduct indexes (hattach cpu (hproduct standardDTypes shapes)))           Device {deviceType = CUDA, deviceIndex = 0} -> do             hfoldrM @IO DiagflatSpec () (hproduct indexes (hattach cuda0 (hproduct (withHalf standardDTypes) shapes)))+          Device {deviceType = MPS, deviceIndex = 0} -> do+            hfoldrM @IO DiagflatSpec () (hproduct indexes (hattach mps (hproduct (withHalf mpsDTypes) shapes)))       it "diagonal" $ do         let shapes1 = Proxy @'[2, 5, 4, 2] :. HNil             shapes2 = Proxy @'[2, 3] :. shapes1@@ -314,3 +331,7 @@             hfoldrM @IO DiagonalSpec () (hproduct (hproduct tris (hproduct indexes dims)) (hattach cuda0 (hproduct (withHalf standardDTypes) allShapes)))             hfoldrM @IO DiagonalSpec () (hproduct (hproduct tris (hproduct allIndexes allDims)) (hattach cuda0 (hproduct (withHalf standardDTypes) shapes1)))             hfoldrM @IO DiagonalSpec () (hproduct (hproduct tris (hproduct allIndexes dims)) (hattach cuda0 (hproduct (withHalf standardDTypes) shapes2)))+          Device {deviceType = MPS, deviceIndex = 0} -> do+            hfoldrM @IO DiagonalSpec () (hproduct (hproduct tris (hproduct indexes dims)) (hattach mps (hproduct (withHalf mpsDTypes) allShapes)))+            hfoldrM @IO DiagonalSpec () (hproduct (hproduct tris (hproduct allIndexes allDims)) (hattach mps (hproduct (withHalf mpsDTypes) shapes1)))+            hfoldrM @IO DiagonalSpec () (hproduct (hproduct tris (hproduct allIndexes dims)) (hattach mps (hproduct (withHalf mpsDTypes) shapes2)))

@@ -540,6 +540,8 @@               hfoldrM @IO lossSpec () (hproduct reductions (hattach cpu (hproduct standardFloatingPointDTypes standardShapes)))             Device {deviceType = CUDA, deviceIndex = 0} ->               hfoldrM @IO lossSpec () (hproduct reductions (hattach cuda0 (hproduct allFloatingPointDTypes standardShapes)))+            Device {deviceType = MPS, deviceIndex = 0} ->+              hfoldrM @IO lossSpec () (hproduct reductions (hattach mps (hproduct mpsFloatingPointDTypes standardShapes)))       it "binaryCrossEntropy" $ dispatch BinaryCrossEntropySpec       it "mseLoss" $ dispatch MSELossSpec @@ -551,6 +553,8 @@               hfoldrM @IO softmaxSpec () (hproduct softmaxDims (hattach cpu (hproduct standardFloatingPointDTypes standardShapes)))             Device {deviceType = CUDA, deviceIndex = 0} ->               hfoldrM @IO softmaxSpec () (hproduct softmaxDims (hattach cuda0 (hproduct allFloatingPointDTypes standardShapes)))+            Device {deviceType = MPS, deviceIndex = 0} ->+              hfoldrM @IO softmaxSpec () (hproduct softmaxDims (hattach mps (hproduct mpsFloatingPointDTypes standardShapes)))       it "softmax" $ dispatch SoftmaxSpec       it "logSoftmax" $ dispatch LogSoftmaxSpec @@ -560,16 +564,22 @@           hfoldrM @IO DotSpec () (hattach cpu (hproduct standardDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO DotSpec () (hattach cuda0 (hproduct allFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO DotSpec () (hattach mps (hproduct mpsFloatingPointDTypes (Proxy @1 :. Proxy @2 :. HNil)))       it "inverse" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO InverseSpec () (hattach cpu (hproduct standardFloatingPointDTypes squareShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO InverseSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes (Proxy @'[1, 1] :. Proxy @'[2, 2] :. Proxy @'[1, 1, 1] :. Proxy @'[2, 2, 2] :. HNil)))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO InverseSpec () (hattach mps (hproduct mpsFloatingPointDTypes (Proxy @'[1, 1] :. Proxy @'[2, 2] :. Proxy @'[1, 1, 1] :. Proxy @'[2, 2, 2] :. HNil)))       let dispatchSymeigSpec symeigSpec = case device of             Device {deviceType = CPU, deviceIndex = 0} ->               hfoldrM @IO symeigSpec () (hattach cpu (hproduct standardFloatingPointDTypes squareShapes))             Device {deviceType = CUDA, deviceIndex = 0} ->               hfoldrM @IO symeigSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes squareShapes))+            Device {deviceType = MPS, deviceIndex = 0} ->+              hfoldrM @IO symeigSpec () (hattach mps (hproduct mpsFloatingPointDTypes squareShapes))       it "symeig" $ do         dispatchSymeigSpec SymeigSpec       it "symeigvalues" $ do@@ -582,6 +592,8 @@             hfoldrM @IO EigSpec () (hproduct eigenVectors (hattach cpu (hproduct standardFloatingPointDTypes ns)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO EigSpec () (hproduct eigenVectors (hattach cuda0 (hproduct standardFloatingPointDTypes ns)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO EigSpec () (hproduct eigenVectors (hattach mps (hproduct mpsFloatingPointDTypes ns)))       it "svd" $ do         let svdShapes = Proxy @'[1, 1] :. Proxy @'[1, 2] :. Proxy @'[2, 1] :. Proxy @'[1, 1, 1] :. Proxy @'[3, 2, 3] :. Proxy @'[3, 3, 2] :. HNil             reducedSVD = Proxy @'ThinSVD :. Proxy @'FullSVD :. HNil@@ -590,11 +602,15 @@             hfoldrM @IO SVDSpec () (hproduct reducedSVD (hattach cpu (hproduct standardFloatingPointDTypes svdShapes)))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO SVDSpec () (hproduct reducedSVD (hattach cuda0 (hproduct standardFloatingPointDTypes svdShapes)))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO SVDSpec () (hproduct reducedSVD (hattach mps (hproduct mpsFloatingPointDTypes svdShapes)))       it "cholesky" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO CholeskySpec () (hattach cpu (hproduct standardFloatingPointDTypes squareShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO CholeskySpec () (hattach cuda0 (hproduct standardFloatingPointDTypes squareShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO CholeskySpec () (hattach mps (hproduct mpsFloatingPointDTypes squareShapes))       it "choleskyInverse" $ do         let choleskyInverseShapes = Proxy @'[1, 1] :. Proxy @'[2, 2] :. HNil         case device of@@ -602,6 +618,8 @@             hfoldrM @IO CholeskyInverseSpec () (hattach cpu (hproduct standardFloatingPointDTypes choleskyInverseShapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO CholeskyInverseSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes choleskyInverseShapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO CholeskyInverseSpec () (hattach mps (hproduct mpsFloatingPointDTypes choleskyInverseShapes))       it "choleskySolve" $ do         let choleskySolveShapes =               hzip@@ -612,6 +630,8 @@             hfoldrM @IO CholeskySolveSpec () (hattach cpu (hproduct standardFloatingPointDTypes choleskySolveShapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO CholeskySolveSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes choleskySolveShapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO CholeskySolveSpec () (hattach mps (hproduct mpsFloatingPointDTypes choleskySolveShapes))       it "solve" $ do         let solveShapes =               hzip@@ -622,6 +642,8 @@             hfoldrM @IO SolveSpec () (hattach cpu (hproduct standardFloatingPointDTypes solveShapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO SolveSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes solveShapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO SolveSpec () (hattach mps (hproduct mpsFloatingPointDTypes solveShapes))      describe "boolean algebra" $ do       do@@ -630,6 +652,8 @@                 hfoldrM @IO anyAllSpec () (hattach cpu standardShapes)               Device {deviceType = CUDA, deviceIndex = 0} ->                 hfoldrM @IO anyAllSpec () (hattach cuda0 standardShapes)+              Device {deviceType = MPS, deviceIndex = 0} ->+                hfoldrM @IO anyAllSpec () (hattach mps standardShapes)         it "all" $ dispatch AllSpec         it "any" $ dispatch AnySpec       do@@ -653,6 +677,14 @@                       (hproduct anyPrimeAllPrimeDims keepOrDropDims)                       (hattach cuda0 anyPrimeAllPrimeShapes)                   )+              Device {deviceType = MPS, deviceIndex = 0} ->+                hfoldrM @IO+                  anyPrimeAllPrimeSpec+                  ()+                  ( hproduct+                      (hproduct anyPrimeAllPrimeDims keepOrDropDims)+                      (hattach mps anyPrimeAllPrimeShapes)+                  )         it "allDim" $ dispatch AllPrimeSpec         it "anyDim" $ dispatch AnyPrimeSpec @@ -692,6 +724,8 @@             hfoldrM @IO LstmCellSpec () (hattach cpu (hproduct standardFloatingPointDTypes sizes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO LstmCellSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes sizes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO LstmCellSpec () (hattach mps (hproduct mpsFloatingPointDTypes sizes))       it "gruCell op" $ do         let sizes =               hzip3@@ -703,3 +737,5 @@             hfoldrM @IO GruCellSpec () (hattach cpu (hproduct standardFloatingPointDTypes sizes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO GruCellSpec () (hattach cuda0 (hproduct standardFloatingPointDTypes sizes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO GruCellSpec () (hattach mps (hproduct mpsFloatingPointDTypes sizes))

@@ -398,45 +398,66 @@         hfoldrM @IO GDConvQuadSpec () (hattach cpu (hproduct standardFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO GDConvQuadSpec () (hattach cuda0 (hproduct allFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO GDConvQuadSpec () (hattach mps (hproduct mpsFloatingPointDTypes (Proxy @1 :. Proxy @2 :. HNil)))     it "Rosenbrock" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO GDRosenbrockSpec () (hattach cpu standardFloatingPointDTypes)       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO GDRosenbrockSpec () (hattach cuda0 standardFloatingPointDTypes)+      Device {deviceType = MPS, deviceIndex = 0} ->+        return ()+      -- ToDo: This test does not pass. --+      -- Device {deviceType = MPS, deviceIndex = 0} ->+      --   hfoldrM @IO GDRosenbrockSpec () (hattach mps mpsFloatingPointDTypes)     it "Ackley" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO GDAckleySpec () (hattach cpu standardFloatingPointDTypes)       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO GDAckleySpec () (hattach cuda0 standardFloatingPointDTypes)+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO GDAckleySpec () (hattach mps mpsFloatingPointDTypes)   describe "GDM" $ do     it "convex quadratic" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO GDMConvQuadSpec () (hattach cpu (hproduct standardFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO GDMConvQuadSpec () (hattach cuda0 (hproduct allFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO GDMConvQuadSpec () (hattach mps (hproduct mpsFloatingPointDTypes (Proxy @1 :. Proxy @2 :. HNil)))     it "Rosenbrock" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO GDMRosenbrockSpec () (hattach cpu standardFloatingPointDTypes)       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO GDMRosenbrockSpec () (hattach cuda0 standardFloatingPointDTypes)+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO GDMRosenbrockSpec () (hattach mps mpsFloatingPointDTypes)     it "Ackley" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO GDMAckleySpec () (hattach cpu standardFloatingPointDTypes)       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO GDMAckleySpec () (hattach cuda0 standardFloatingPointDTypes)+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO GDMAckleySpec () (hattach mps mpsFloatingPointDTypes)   describe "Adam" $ do     it "convex quadratic" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO AdamConvQuadSpec () (hattach cpu (hproduct standardFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO AdamConvQuadSpec () (hattach cuda0 (hproduct allFloatingPointDTypes (Proxy @0 :. Proxy @1 :. Proxy @2 :. HNil)))+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO AdamConvQuadSpec () (hattach mps (hproduct mpsFloatingPointDTypes (Proxy @1 :. Proxy @2 :. HNil)))     it "Rosenbrock" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO AdamRosenbrockSpec () (hattach cpu standardFloatingPointDTypes)       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO AdamRosenbrockSpec () (hattach cuda0 standardFloatingPointDTypes)+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO AdamRosenbrockSpec () (hattach mps mpsFloatingPointDTypes)     it "Ackley" $ case device of       Device {deviceType = CPU, deviceIndex = 0} ->         hfoldrM @IO AdamAckleySpec () (hattach cpu standardFloatingPointDTypes)       Device {deviceType = CUDA, deviceIndex = 0} ->         hfoldrM @IO AdamAckleySpec () (hattach cuda0 standardFloatingPointDTypes)+      Device {deviceType = MPS, deviceIndex = 0} ->+        hfoldrM @IO AdamAckleySpec () (hattach mps mpsFloatingPointDTypes)

@@ -85,6 +85,7 @@         broadcastableShapes1 = Proxy @'[2, 1, 1] :. HNil         standardDTypes2 = hproduct standardDTypes standardDTypes         almostAllDTypes2 = hproduct (withHalf standardDTypes) (withHalf standardDTypes)+        mpsDTypes2 = hproduct mpsDTypes mpsDTypes         identicalShapes = hzip standardShapes standardShapes         broadcastableShapes = hzip broadcastableShapes0 broadcastableShapes1 @@ -96,12 +97,16 @@                   hfoldrM @IO binarySpec () (hattach cpu (hproduct standardDTypes2 identicalShapes))                 Device {deviceType = CUDA, deviceIndex = 0} ->                   hfoldrM @IO binarySpec () (hattach cuda0 (hproduct almostAllDTypes2 identicalShapes))+                Device {deviceType = MPS, deviceIndex = 0} ->+                  hfoldrM @IO binarySpec () (hattach mps (hproduct mpsDTypes2 identicalShapes))             it "works on broadcastable tensors of different shapes" $               case device of                 Device {deviceType = CPU, deviceIndex = 0} ->                   hfoldrM @IO binarySpec () (hattach cpu (hproduct standardDTypes2 broadcastableShapes))                 Device {deviceType = CUDA, deviceIndex = 0} ->                   hfoldrM @IO binarySpec () (hattach cuda0 (hproduct almostAllDTypes2 broadcastableShapes))+                Device {deviceType = MPS, deviceIndex = 0} ->+                  hfoldrM @IO binarySpec () (hattach mps (hproduct mpsDTypes2 broadcastableShapes))       describe "addition" $ dispatch AddSpec       describe "subtraction" $ dispatch SubSpec       describe "multiplication" $ dispatch MulSpec@@ -114,6 +119,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "returns the matrix-matrix product if both arguments are 2-dimensional" $ do         let shapes = hzip (Proxy @'[3, 2] :. HNil) (Proxy @'[2, 4] :. HNil)         case device of@@ -121,6 +128,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "returns the matrix-matrix product if the first argument is 1-dimensional and the second argument is 2-dimensional by temporarily adding a 1 to the dimension of the first argument" $ do         let shapes = hzip (Proxy @'[3] :. HNil) (Proxy @'[3, 4] :. HNil)         case device of@@ -128,6 +137,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "returns the matrix-vector product if the first argument is 2-dimensional and the second argument is 1-dimensional" $ do         let shapes = hzip (Proxy @'[3, 4] :. HNil) (Proxy @'[4] :. HNil)         case device of@@ -135,6 +146,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "returns a batched matrix-matrix product if both arguments are at least 2-dimensional and the batch (i.e. non-matrix) dimensions are broadcastable" $ do         let shapes = hzip (Proxy @'[2, 1, 4, 3] :. HNil) (Proxy @'[3, 3, 2] :. HNil)         case device of@@ -142,6 +155,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "returns a batched matrix-matrix product if the first argument is 1-dimensional and the second argument has more than 2 dimensions" $ do         let shapes = hzip (Proxy @'[3] :. HNil) (Proxy @'[2, 3, 4] :. HNil)         case device of@@ -149,6 +164,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))       it "returns a batched matrix-vector product if the first argument has more than 2 dimensions and the second argument is 1-dimensional" $ do         let shapes = hzip (Proxy @'[2, 3, 4] :. HNil) (Proxy @'[4] :. HNil)         case device of@@ -156,6 +173,8 @@             hfoldrM @IO MatMulSpec () (hattach cpu (hproduct standardDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO MatMulSpec () (hattach cuda0 (hproduct allFloatingPointDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO MatMulSpec () (hattach mps (hproduct mpsFloatingPointDTypes shapes))  testTensorListFold ::   forall device dtype shape. Tensor device dtype shape -> IO [Torch.ATenTensor]

@@ -1,3 +1,4 @@+{-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-}@@ -141,6 +142,7 @@         broadcastableShapes1 = Proxy @'[2, 1, 1] :. HNil         standardDTypes2 = hproduct standardDTypes standardDTypes         almostAllDTypes2 = hproduct (withHalf standardDTypes) (withHalf standardDTypes)+        mpsDTypes2 = hproduct mpsDTypes mpsDTypes         identicalShapes = hzip standardShapes standardShapes         broadcastableShapes = hzip broadcastableShapes0 broadcastableShapes1 @@ -152,12 +154,16 @@                   hfoldrM @IO binaryCmpSpec () (hattach cpu (hproduct standardDTypes2 identicalShapes))                 Device {deviceType = CUDA, deviceIndex = 0} ->                   hfoldrM @IO binaryCmpSpec () (hattach cuda0 (hproduct almostAllDTypes2 identicalShapes))+                Device {deviceType = MPS, deviceIndex = 0} ->+                  hfoldrM @IO binaryCmpSpec () (hattach mps (hproduct mpsDTypes2 identicalShapes))             it "works on broadcastable tensors of different shapes" $               case device of                 Device {deviceType = CPU, deviceIndex = 0} ->                   hfoldrM @IO binaryCmpSpec () (hattach cpu (hproduct standardDTypes2 broadcastableShapes))                 Device {deviceType = CUDA, deviceIndex = 0} ->                   hfoldrM @IO binaryCmpSpec () (hattach cuda0 (hproduct almostAllDTypes2 broadcastableShapes))+                Device {deviceType = MPS, deviceIndex = 0} ->+                  hfoldrM @IO binaryCmpSpec () (hattach mps (hproduct mpsDTypes2 broadcastableShapes))       describe "greater than" $ dispatch GTSpec       describe "lower than" $ dispatch LTSpec       describe "greater or equal than" $ dispatch GESpec@@ -175,17 +181,27 @@             hfoldrM @IO ReshapeSpec () (hattach cpu (hproduct allDTypes shapes))           Device {deviceType = CUDA, deviceIndex = 0} ->             hfoldrM @IO ReshapeSpec () (hattach cuda0 (hproduct allDTypes shapes))+          Device {deviceType = MPS, deviceIndex = 0} ->+            hfoldrM @IO ReshapeSpec () (hattach mps (hproduct mpsDTypes shapes))        it "toDevice" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->+#ifdef __APPLE__+          hfoldrM @IO ToDeviceSpec () (hattach cpu (hproduct mpsDTypes standardShapes))+#else           hfoldrM @IO ToDeviceSpec () (hattach cpu (hproduct allDTypes standardShapes))+#endif         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO ToDeviceSpec () (hattach cuda0 (hproduct allDTypes standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO ToDeviceSpec () (hattach mps (hproduct mpsDTypes standardShapes))       it "toType" $ case device of         Device {deviceType = CPU, deviceIndex = 0} ->           hfoldrM @IO ToTypeSpec () (hattach cpu (hproduct (hproduct allDTypes allDTypes) standardShapes))         Device {deviceType = CUDA, deviceIndex = 0} ->           hfoldrM @IO ToTypeSpec () (hattach cuda0 (hproduct (hproduct allDTypes allDTypes) standardShapes))+        Device {deviceType = MPS, deviceIndex = 0} ->+          hfoldrM @IO ToTypeSpec () (hattach mps (hproduct (hproduct mpsDTypes mpsDTypes) standardShapes))      describe "untyped to typed tensor" $ do       it "withTensor" $ do