Experiment with differentiation via typeclasses

bench/Bench.hs

@@ -4,7 +4,7 @@
 
 import MLambda.Matrix
 import MLambda.NDArr
 import MLambda.TypeLits (KnownNat, natVal)
 
 import Data.Massiv.Array (Array, Comp (..), Ix2, pattern Sz2)
 import Data.Massiv.Array.Manifest (S)
@@ -13,7 +13,7 @@
 import Data.Random.Normal (normalIO)
 import Data.Vector.Storable qualified as Storable
 import Foreign.Storable
-import GHC.TypeLits (type (*), type (<=))
+import GHC.TypeLits hiding (natVal)
 import System.Random (mkStdGen, setStdGen)
 import Test.Tasty (localOption)
 import Test.Tasty.Bench
@@ -32,7 +32,7 @@
 
 mkVec :: forall m n -> (KnownNat m, KnownNat n, Storable a)
       => IO a -> IO (Storable.Vector a)
-mkVec m n gen = Storable.replicateM (natVal n * natVal m) $ gen
+mkVec m n = Storable.replicateM (natVal n * natVal m)
 
 mkMassiv :: forall m n -> (KnownNat m, KnownNat n, Storable a) => IO a -> IO (Array S Ix2 a)
 mkMassiv m n gen = freeze Seq =<< makeMArrayS (Sz2 (natVal n) (natVal m)) (const gen)

ml.cabal

@@ -1,6 +1,6 @@
 cabal-version: 2.2
 
--- This file has been generated from package.yaml by hpack version 0.38.1.
+-- This file has been generated from package.yaml by hpack version 0.39.6.
 --
 -- see: https://github.com/sol/hpack
 
@@ -25,6 +25,7 @@ source-repository head
 
 library
   exposed-modules:
+      MLambda.Differentiable
       MLambda.Index
       MLambda.Linear
       MLambda.Matrix
@@ -55,6 +56,7 @@ library
     , singletons-base
     , template-haskell
     , vector
+    , vinyl
   default-language: GHC2024
 
 test-suite ml-test
@@ -81,7 +83,7 @@ test-suite ml-test
       base >=4.7 && <5
     , blas-ffi
     , deepseq
-    , falsify
+    , falsify >=0.4.0 && <0.5
     , haskell-src-meta
     , massiv
     , ml
@@ -91,9 +93,11 @@ test-suite ml-test
     , singletons
     , singletons-base
     , tasty
+    , tasty-falsify
     , tasty-hunit
     , template-haskell
     , vector
+    , vinyl
   default-language: GHC2024
 
 benchmark ml-bench
@@ -131,4 +135,5 @@ benchmark ml-bench
     , tasty-bench
     , template-haskell
     , vector
+    , vinyl
   default-language: GHC2024

package.yaml

@@ -32,6 +32,7 @@ dependencies:
 - singletons
 - singletons-base
 - primitive
+- vinyl
 # - ghc-typelits-natnormalise
 
 ghc-options:
@@ -77,7 +78,8 @@ tests:
     - ml
     - tasty
     - tasty-hunit
-    - falsify
+    - falsify >= 0.4.0 && < 0.5
+    - tasty-falsify
 
 language: GHC2024
 default-extensions:

src/MLambda/Differentiable.hs

@@ -0,0 +1,110 @@
+-- |
+-- Module      : MLambda.Differentiable
+-- Description : Implements differentiation machinery.
+-- Copyright   : (c) neclitoris, 2026
+-- License     : BSD-3-Clause
+-- Maintainer  : nas140301@gmail.com
+-- Stability   : experimental
+-- Portability : portable
+--
+-- This module contains definition of `Differentiable` type class.
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE RequiredTypeArguments #-}
+module MLambda.Differentiable
+  ( Functional(..)
+  , Differentiable(..)
+  , Matmul(..)
+  ) where
+
+import MLambda.Index
+import MLambda.Matrix
+import MLambda.NDArr
+import MLambda.TypeLits
+
+import Data.Bifunctor
+import Data.Either.Singletons
+import Data.Kind
+import Data.List.Singletons
+import Data.Singletons
+import Data.Type.Equality
+import Data.Vinyl hiding ((:~:))
+import Numeric.Netlib.Class
+
+import Unsafe.Coerce
+
+import Prelude hiding (Floating)
+
+
+type family ArgsL (i :: [[Natural]]) e :: [Type] where
+  ArgsL '[] e = '[]
+  ArgsL (x ': xs) e = NDArr x e : ArgsL xs e
+
+type Args i e = Rec At (Fins (ArgsL i e))
+
+type Fun i o e = Args i e -> NDArr o e
+
+class Functional f i o e where
+  ($$) :: f -> Fun i o e
+
+class Functional f i o e => Differentiable f i o e where
+  d :: f -> Args i e -> Args (Map (Apply (++@#@$) o) i) e
+
+data Matmul = Matmul
+
+instance (KnownNat m, KnownNat n, KnownNat k, Floating e) => Functional Matmul '[[m,n], [n,k]] '[m,k] e where
+  _ $$ (At a :& At b :& RNil) = a `cross` b
+
+instance (1 <= m, 1 <= n, 1 <= k, KnownNat m, KnownNat n, KnownNat k, Floating e) => Differentiable Matmul '[[m,n], [n,k]] '[m,k] e where
+  d _ (At a :& At b :& RNil) = At a' :& At b' :& RNil
+    where -- TODO: optimize both representation and performance of this
+      a' = fromIndex \(i :. j :. k :. l) -> if k == i then b `at` (l :. j) else 0
+      b' = fromIndex \(i :. j :. k :. l) -> if l == j then a `at` (i :. k) else 0
+
+data (:.:) f g = f :.: g
+
+splitRec :: forall {l :: [Type]} (l1 :: [Type]) (l2 :: [Type]) .
+  (l ~ l1 ++ l2) => Sing l1 -> Sing l -> Rec At (Fins (l1 ++ l2)) -> (Rec At (Fins l1), Rec At (Fins l2))
+splitRec SNil _ xs = (RNil, xs)
+splitRec (SCons _ (sxs :: Sing xs)) (SCons _ ys) (At v :& r) =
+  first (\lhs -> At v :& shiftFS (withSingI sxs singFins) lhs) $ splitRec @xs @l2 sxs ys (stripFS (withSingI ys singFins) r)
+
+{-# NOINLINE[1] unFS #-}
+{-# RULES "unFSnop" unFS = unsafeCoerce #-}
+unFS :: At (FS i) -> At i
+unFS (At x) = At x
+
+{-# NOINLINE[1] doFS #-}
+{-# RULES "doFSnop" doFS = unsafeCoerce #-}
+doFS :: At i -> At (FS i)
+doFS (At x) = At x
+
+{-# NOINLINE[1] stripFS #-}
+{-# RULES "stripFSnop" forall x . stripFS x = unsafeCoerce #-}
+stripFS
+  :: Sing is
+  -> Rec At (Map (TyCon1 FS) is)
+  -> Rec At is
+stripFS SNil RNil               = RNil
+stripFS (SCons _ sis) (x :& xs) = unFS x :& stripFS sis xs
+
+{-# NOINLINE[1] shiftFS #-}
+{-# RULES "shiftFSnop" forall x . shiftFS x = unsafeCoerce #-}
+shiftFS
+  :: Sing is
+  -> Rec At is
+  -> Rec At (Map (TyCon1 FS) is)
+shiftFS SNil RNil               = RNil
+shiftFS (SCons _ sis) (x :& xs) = doFS x :& shiftFS sis xs
+
+instance
+    ( Functional f1 i1 o1 e
+    , Functional f2 (o1 : i2) o2 e
+    , ArgsL i e ~ ArgsL i1 e ++ ArgsL i2 e, SingI (ArgsL i1 e), SingI (ArgsL i2 e))
+  => Functional (f2 :.: f1) i o2 e where
+  (f2 :.: f1) $$ r =
+    let sl1 = sing @(ArgsL i1 e)
+        sl2 = sing @(ArgsL i2 e)
+        sl = sl1 %++ sl2
+        (lhs, rhs) = splitRec @(ArgsL i1 e) @(ArgsL i2 e) sl1 sl r
+     in ($$) @f2 @(o1 : i2) @o2 @e f2 (At (($$) @f1 @i1 @o1 @e f1 lhs) :& shiftFS (withSingI sl2 singFins) rhs)
+

src/MLambda/TypeLits.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE RequiredTypeArguments #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
 
 -- |
 -- Module      : MLambda.TypeLits
@@ -22,12 +23,20 @@ module MLambda.TypeLits
   , Peano
   , RNat (..)
   , RPNat (..)
+  , Fin (..)
+  , SFin (..)
+  , SingFins(..)
+  , Fins
+  , At(..)
+  , type (!)
   , ReifiedNat
   , rnat
   , rpnat
   ) where
 
-import Data.Proxy (Proxy (Proxy))
+import Data.Kind
+import Data.List.Singletons (Map, SList (..), sMap, type (++))
+import Data.Singletons
 import GHC.TypeError (ErrorMessage (..), TypeError)
 import GHC.TypeNats hiding (natVal)
 
@@ -67,6 +76,63 @@ data RPNat n where
   RPZ :: RPNat PZ
   RPS :: RPNat n -> RPNat (PS n)
 
+-- | Finite list index.
+type Fin :: [k] -> Type
+data Fin (l :: [k]) where
+  FZ :: Fin (x : xs)
+  FS :: Fin xs -> Fin (x : xs)
+
+type SFin :: forall k (l :: [k]) . Fin l -> Type
+data SFin (f :: Fin l) where
+  SFZ :: SFin FZ
+  SFS :: forall xs (f :: Fin xs) . SFin f -> SFin (FS f)
+
+type instance Sing = SFin
+
+instance SingKind (Fin l) where
+  type Demote (Fin l) = Fin l
+
+  fromSing SFZ     = FZ
+  fromSing (SFS s) = FS $ fromSing s
+
+  toSing FZ     = SomeSing SFZ
+  toSing (FS s) = (\(SomeSing s') -> SomeSing $ SFS s') $ toSing s
+
+type (!) :: forall k . forall (l :: [k]) -> Fin l -> k
+type family (!) (l :: [k]) (i :: Fin l) where
+  (x ': _) ! FZ = x
+  (_ ': xs) ! (FS i) = xs ! i
+
+type At :: forall {l :: [Type]} . Fin l -> Type
+data At (i :: Fin l) where
+  At :: forall {l} (i :: Fin l) . l ! i -> At i
+
+type AppendFin :: forall r -> Fin l -> Fin (l ++ r)
+type family AppendFin r f where
+  AppendFin r FZ = FZ
+  AppendFin r (FS s) = FS (AppendFin r s)
+
+type PrependFin :: forall l -> Fin r -> Fin (l ++ r)
+type family PrependFin l f where
+  PrependFin '[] s = s
+  PrependFin (x:xs) s = FS (PrependFin xs s)
+
+-- | Creates a list of indices into a type-level list.
+type Fins :: forall {k} . forall (l :: [k]) -> [Fin l]
+type family Fins (l :: [k :: Type]) where
+  Fins '[] = '[]
+  Fins (x ': xs) = FZ ': Map (TyCon1 FS) (Fins xs)
+
+class SingFins l where
+  singFins :: Sing (Fins l)
+
+instance SingI l => SingFins l where
+  singFins =
+    case sing @l of
+      SNil -> SNil
+      SCons (_ :: Sing x) (sxs :: Sing xs) -> withSingI sxs $
+        SCons SFZ (sMap @(Fin xs) @(Fin (x:xs)) @(TyCon1 FS) (SLambda SFS) (singFins @xs))
+
 -- | A stronger variant of 'KnownNat' which enables induction on type-level naturals.
 class ReifiedNat n where
   rnat0 :: RNat n

stack.yaml

@@ -17,7 +17,7 @@
 #
 # snapshot: ./custom-snapshot.yaml
 # snapshot: https://example.com/snapshots/2024-01-01.yaml
-snapshot: nightly-2025-08-02
+snapshot: nightly-2026-09-01
 
 # User packages to be built.
 # Various formats can be used as shown in the example below.
@@ -35,7 +35,7 @@ packages:
 # forks / in-progress versions pinned to a git hash. For example:
 #
 extra-deps:
-- falsify-0.2.0@sha256:af5c4142095d05775236c8e18e827d540ed22f57b3b74b2268b32040b514ac88,5451
+- tasty-falsify-0.1.0
 #
 # extra-deps: []
 
@@ -62,6 +62,3 @@ extra-deps:
 #
 # Allow a newer minor version of GHC than the snapshot specifies
 # compiler-check: newer-minor
-allow-newer-deps:
-- falsify
-allow-newer: true

test/Test/MLambda/Matrix.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeAbstractions #-}
 module Test.MLambda.Matrix (testMatrix) where
 
@@ -11,7 +12,8 @@
 import Data.Bool.Singletons
 import Data.Singletons
 import GHC.TypeLits.Singletons
+import Test.Falsify
 import Test.Falsify.Predicate ((.$))
 import Test.Falsify.Predicate qualified as Pred
 import Test.Tasty
 import Test.Tasty.Falsify

test/Test/MLambda/NDArr.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeAbstractions #-}
 {-# LANGUAGE ViewPatterns #-}
 module Test.MLambda.NDArr (testNDArr) where
@@ -15,8 +16,9 @@
 import Data.Singletons
 import GHC.TypeLits.Singletons
 import Prelude.Singletons ((%+))
+import Test.Falsify
 import Test.Falsify.Generator qualified as Gen
 import Test.Falsify.Predicate ((.$))
 import Test.Falsify.Predicate qualified as Pred
 import Test.Tasty
 import Test.Tasty.Falsify

test/Test/MLambda/Utils.hs

@@ -1,28 +1,31 @@
 {-# LANGUAGE TypeAbstractions #-}
 module Test.MLambda.Utils where
 
 import MLambda.Index
 import MLambda.NDArr
 import MLambda.TypeLits
 
+import Data.Falsify.ConcreteFun as D
+import Data.Falsify.ProperFraction as D
 import Data.Proxy
 import Foreign.Storable
+import Test.Falsify
 import Test.Falsify.Generator (Gen)
 import Test.Falsify.Generator qualified as Gen
 import Test.Falsify.Range qualified as Range
 
 genSz :: Gen Natural
-genSz = fromIntegral <$> Gen.int (Range.between (1, 5))
+genSz = fromIntegral <$> Gen.int (Range.inclusive (1, 5))
 
 genDim :: Word -> Word -> Gen [Natural]
-genDim a b = Gen.list (Range.between (a, b)) genSz
+genDim a b = Gen.list (Range.inclusive (a, b)) genSz
 
 genInt :: Gen Int
-genInt = Gen.inRange $ Range.between (-1000, 1000)
+genInt = Gen.inRange $ Range.inclusive (-1000, 1000)
 
 genDouble :: Gen Double
 genDouble = Gen.inRange $ Range.fromProperFraction 64
-  \(Range.ProperFraction d) -> 1 + 4 * d
+  \(D.ProperFraction d) -> 1 + 4 * d
 
 genIndex :: forall dim . Ix dim => Gen (Index dim)
 genIndex = case inst @dim of
@@ -34,9 +37,9 @@
 
 genNDArr :: forall dim e . (Storable e, Ix dim) => Gen e -> Gen (NDArr dim e)
 genNDArr g = do
-  Gen.Fn f <- Gen.fun g
+  Fn f <- Gen.fun g
   pure $ fromIndex f
 
 instance Enum (Index dim) => Gen.Function (Index dim) where
-  function gb = Gen.functionMap fromEnum toEnum <$> Gen.function gb
+  function gb = D.map fromEnum toEnum <$> Gen.function gb