Experiment with differentiation via typeclasses

ml.cabal

@@ -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
@@ -94,6 +96,7 @@ test-suite ml-test
     , tasty-hunit
     , template-haskell
     , vector
+    , vinyl
   default-language: GHC2024
 
 benchmark ml-bench
@@ -131,4 +134,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:

src/MLambda/Differentiable.hs

@@ -0,0 +1,63 @@
+-- |
+-- 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 PatternSynonyms #-}
+module MLambda.Differentiable
+  ( Functional(..)
+  , Differentiable(..)
+  , Matmul(..)
+  ) where
+
+import MLambda.Index
+import MLambda.Matrix
+import MLambda.NDArr
+import MLambda.TypeLits
+
+import Data.Kind
+import Data.List.Singletons
+import Data.Vinyl
+import Numeric.Netlib.Class
+
+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 family Args (i :: [[Natural]]) e :: Type where
+  Args i e = Rec At (Fins @(ArgsL i e) (ArgsL i e))
+
+type family Fun (i :: [[Natural]]) (o :: [Natural]) e :: Type where
+  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 -> Index o -> Args 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) (i :. j) = At a' :& At b' :& RNil
+    where
+      a' = fromIndex \(k :. l) -> if k == i then b `at` (l :. j) else 0
+      b' = fromIndex \(k :. l) -> if l == j then a `at` (i :. k) else 0
+
+data (:.:) f1 f2 = f1 :.: f2
+
+instance (Functional f1 i1 o1 e, Functional f2 (o1 : i2) o2 e, i ~ i1 ++ i2) => Functional (f1 :.: f2) i o2 e where
+  (f1 :.: f2) $$ r = undefined
+

src/MLambda/TypeLits.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE RequiredTypeArguments #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
 
 -- |
 -- Module      : MLambda.TypeLits
@@ -22,11 +23,16 @@ module MLambda.TypeLits
   , Peano
   , RNat (..)
   , RPNat (..)
+  , Fin (..)
+  , Fins
+  , At(..)
+  , type (!)
   , ReifiedNat
   , rnat
   , rpnat
   ) where
 
+import Data.Kind
 import Data.Proxy (Proxy (Proxy))
 import GHC.TypeError (ErrorMessage (..), TypeError)
 import GHC.TypeNats hiding (natVal)
@@ -67,6 +73,32 @@ data RPNat n where
   RPZ :: RPNat PZ
   RPS :: RPNat n -> RPNat (PS n)
 
+-- | Finite list index.
+type Fin :: [k] -> Type
+data Fin (l :: [k :: Type]) where
+  FZ :: Fin (x : xs)
+  FS :: Fin xs -> Fin (x : xs)
+
+type (!) :: forall l -> Fin l -> Type
+type family (!) (l :: [k :: Type]) (i :: Fin l) where
+  (x ': _) ! FZ = x
+  (_ ': xs) ! (FS i) = xs ! i
+
+type At :: forall {k} {l :: [k]} . Fin l -> Type
+data At (i :: Fin l) where
+  At :: l ! i -> At i
+
+type Map :: forall k l . (k -> l) -> [k] -> [l]
+type family Map f l where
+  Map _ '[] = '[]
+  Map f (x ': xs) = f x ': Map f xs
+
+-- | Creates a list of indices into a type-level list.
+type Fins :: forall {k} (l :: [k]) . [k] -> [Fin l]
+type family Fins (l :: [k :: Type]) where
+  Fins '[] = '[]
+  Fins (x ': xs) = FZ ': Map FS (Fins xs)
+
 -- | A stronger variant of 'KnownNat' which enables induction on type-level naturals.
 class ReifiedNat n where
   rnat0 :: RNat n