0 | module NN.Architectures.Affine
 1 |
 2 | import System.Random
 3 |
 4 | import Data.Tensor
 5 | import Data.Para
 6 | import Data.ComMonoid
 7 | import Data.Container.Additive
 8 | import Data.Autodiff.Model
 9 |
10 | -- This is often called a 'linear layer', but really it is affine because of the bias
11 |
12 | ||| A weight matrix and a bias vector
13 | public export
14 | AffineParams : (x, y : Axis) -> y `ConsistentWith` [x] =>
15 |   Type -> Type
16 | AffineParams x y a = (Tensor [y, x] a, Tensor [y] a)
17 |
18 | public export
19 | affineImpl : {x, y : Axis} ->
20 |   y `ConsistentWith` [x] =>
21 |   Num a =>
22 |   AllAlgebra [x] a =>
23 |   TensorMonoid x.cont => TensorMonoid y.cont =>
24 |   DPair (Tensor [x] a) (const (AffineParams x y a)) -> Tensor [y] a
25 | affineImpl (input ** (weights, bias))
26 |   = matrixVectorProduct weights input + bias
27 |
28 | public export
29 | affinePara : {x, y : Axis} -> {a : Type} -> Num a =>
30 |   y `ConsistentWith` [x] =>
31 |   AllAlgebra [x] a =>
32 |   TensorMonoid x.cont => TensorMonoid y.cont =>
33 |   Tensor [x] a -\-> Tensor [y] a
34 | affinePara = MkPara
35 |   (const (AffineParams x y a))
36 |   affineImpl
37 |
38 | public export
39 | affineModel : {x, y : Axis} -> {a : Type} ->
40 |   Neg a =>
41 |   y `ConsistentWith` [x] =>
42 |   AllAlgebra [x] a =>
43 |   TensorMonoid x.cont => TensorMonoid y.cont =>
44 |   Algebra (Ext y.cont) (Tensor [x] a) =>
45 |   Random (Tensor [y, x] a) => Random (Tensor [y] a) =>
46 |   Const (Tensor [x] a) -\-> Const (Tensor [y] a)
47 | affineModel = layer (AffineParams x y a)
48 |   [| MkPair (randomRIO (-1, 1)) (randomRIO (-1, 1)) |]
49 |   (\input, (weights, bias) =>
50 |     (matrixVectorProduct weights input + bias ** \dy =>
51 |       (vectorMatrixProduct dy weights, (outer dy input, dy))))
52 |