Add splitLayers

Author: junjihashimoto

src/Torch/Compose.hs

@@ -20,6 +20,8 @@
 {-# LANGUAGE TypeFamilyDependencies#-}
 {-# LANGUAGE TypeOperators#-}
 {-# LANGUAGE UndecidableInstances#-}
+{-# LANGUAGE AllowAmbiguousTypes#-}
+
 
 module Torch.Compose where
 
@@ -31,6 +33,7 @@ import Data.Coerce
 import Control.Exception
 import System.IO.Unsafe
 import qualified Language.Haskell.TH as TH
+import GHC.TypeLits
 -- import qualified Language.Haskell.TH.Syntax as TH
 
 
@@ -144,15 +147,34 @@ instance (HasForwardAssoc f a) => HasForwardAssoc f (HList '[a]) where
   forwardAssoc f (HCons a HNil) = toHList $ forwardAssoc f a
   forwardStochAssoc f (HCons a HNil) = toHList <$> forwardStochAssoc f a
 
+type family ToHNat (x :: Nat) :: HNat where
+  ToHNat 0 = HZero
+  ToHNat x = HSucc (ToHNat ( x - 1 ))
+
+dropLayers :: forall num a ys xs. (KnownNat num, Coercible a (HList xs), HDrop (ToHNat num) xs ys,  HLengthGe xs (ToHNat num)) => a -> HList ys
+dropLayers m = hDrop (Proxy :: Proxy (ToHNat num)) (coerce m)
+
+takeLayers :: forall num a ys xs. (KnownNat num, Coercible a (HList xs), HTake (ToHNat num) xs ys,  SameLength' (HReplicateR (ToHNat num) ()) ys,  HLengthEq1 ys (ToHNat num), HLengthEq2 ys (ToHNat num),  HLengthGe xs (ToHNat num)) => a -> HList ys
+takeLayers m = hTake (Proxy :: Proxy (ToHNat num)) (coerce m)
+
+splitLayers :: forall num a xs ys xsys. (KnownNat num, Coercible a (HList xsys), HSplitAt1 '[] (ToHNat num) xsys xs ys,  HAppendList1 xs ys xsys,  SameLength' (HReplicateR (ToHNat num) ()) xs,  HLengthEq1 xs (ToHNat num), HLengthEq2 xs (ToHNat num)) => a -> (HList xs, HList ys)
+splitLayers m = hSplitAt (Proxy :: Proxy (ToHNat num)) (coerce m)
+
 dropLastLayer :: (Coercible a (HList xs1), HRevApp xs2 '[x] xs1, HRevApp xs2 '[] sx,  HRevApp xs1 '[] (x : xs2), HRevApp sx '[] xs2) => a -> HList sx
 dropLastLayer m = hReverse (hDrop (Proxy :: Proxy (HSucc HZero)) (hReverse (coerce m)))
 
 addLastLayer :: HAppend l1 (HList '[e]) => l1 -> e -> HAppendR l1 (HList '[e])
 addLastLayer a b = a `hAppend` (b .*. HNil)
 
+dropFirstLayer :: Coercible a (HList (x : ys)) => a -> HList ys
+dropFirstLayer m = hDrop (Proxy :: Proxy (HSucc HZero)) (coerce m)
+
 getLastLayer :: (Coercible a (HList l1), HRevApp l1 '[] (e : l)) => a -> e
 getLastLayer a = hLast (coerce a)
 
+getFirstLayer :: Coercible a (HList (e : l)) => a -> e
+getFirstLayer a = hHead (coerce a)
+
 hScanl :: forall f z ls xs1 sx xs2. (HScanr f z ls xs1, HRevApp xs1 '[] sx, HRevApp sx '[] xs1,  HRevApp xs2 '[] ls, HRevApp ls '[] xs2) => f -> z -> HList xs2 -> HList sx
 hScanl a b c = hReverse $ hScanr a b (hReverse c)
 
@@ -259,12 +281,5 @@ instanceForwardAssoc model input output =
   |]
 
 instanceForwardAssocs :: [TH.Q TH.Type] -> TH.Q TH.Type -> TH.Q TH.Type -> TH.DecsQ
-instanceForwardAssocs models input output = do
-  decs <- forM models $ \model ->
-    [d|
-        instance HasForwardAssoc $model $input where
-          type ForwardResult $model $input = $output
-          forwardAssoc = forward
-          forwardStochAssoc = forwardStoch
-    |]
-  return $ concat decs
+instanceForwardAssocs models input output =
+  concat <$> forM models (\model -> instanceForwardAssoc model input output)

stack.yaml

@@ -121,3 +121,30 @@ main = hspec $ do
         zero1 = over (types @Tensor) (ones' . shape) $ getLastLayer m1'
         model' = addLastLayer (dropLastLayer m0) (mergeParameters (+) layer0 zero1)
     return ()
+  it "Fanin" $ do
+    m0 <- sample mlpSpec
+    m1 <- sample mlpSpec
+    let l0 = getFirstLayer m0
+        l1 = getFirstLayer m1
+        fin = Fanin l0 l1
+        model' = HCons fin $ dropFirstLayer m0
+        input = ones' [2,784]
+        out = forward model' (input,input)
+    shape out `shouldBe` [2,10]
+  it "Fanout" $ do
+    m0 <- sample (LinearSpec 10 2)
+    m1 <- sample (LinearSpec 10 3)
+    let fout = Fanout m0 m1
+        input = ones' [1,10]
+        (out0,out1) = forward fout input
+    shape out0 `shouldBe` [1,2]
+    shape out1 `shouldBe` [1,3]
+  it "Split layers" $ do
+    m0 <- sample mlpSpec
+    let (h, t) = splitLayers @2 m0
+        input0 = ones'  [1,784]
+        input1 = ones'  [1,64]
+        output0 = forward h input0
+        output1 = forward t input1
+    shape output0 `shouldBe` [1,64]
+    shape output1 `shouldBe` [1,10]