{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE NoStarIsType #-}

-- | __Differentiable maps__ between products of real vector spaces, with their reverse derivatives.
-- An object is a list of dimensions, @'DF' '[n, m]@ for ℝⁿ × ℝᵐ, and an arrow is a map together
-- with its reverse derivative: given a point and a change of the output, the change of the input
-- it comes from. Running an arrow at a point gives the image together with the reverse derivative
-- there, so the backward pass reuses what the forward pass computed, as in Elliott's
-- /The simple essence of automatic differentiation/. Composition runs the maps forward and chains
-- the reverse derivatives backward, which is the chain rule, so a reverse derivative is never more
-- than a change of vectors passed back through each step.
--
-- The category is cartesian: copying a point is adding the changes that come back, and discarding
-- one is the zero change. Every object is a commutative monoid by adding vectors. 'toLens' gives
-- an arrow as a lens, with the map as the view and the reverse derivative as the update.
--
-- The reverse derivative is a function rather than an arrow of the category itself, so this is the
-- first-order part of a reverse derivative category: the derivative cannot be differentiated again.
module Proarrow.Category.Instance.Diff
  ( DIFF (..)
  , Diff (..)
  , forward
  , backward
  , diff
  , Point (..)
  , pattern Vector
  , coordinates
  , fromCoordinates
  , zero
  , add
  , scale
  , split
  , append
  , linear
  , pointwise
  , times
  , dot
  , matVec
  , matMat
  , broadcast
  , softmax
  , softmaxCrossEntropy
  , bilinear
  , constant
  , toLens
  , Forget
  , AsLens
  ) where

import Data.Coerce (coerce)
import Data.Kind (Type)
import Data.Proxy (Proxy (..))
import Data.Vector.Storable qualified as SV
import GHC.TypeNats (KnownNat, Nat, natVal, type (*))
import Prelude (Double, ($), (*), (+))
import Prelude qualified as P

import Proarrow.Category.Instance.Opposite (OPPOSITE (..), Op (..))
import Proarrow.Category.Instance.TensorNetwork (Size, TNET (..), dimsOf, fromVector, matMul, toVector)
import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), SymMonoidal (..))
import Proarrow.Category.Monoidal.CopyDiscard (CopyDiscard)
import Proarrow.Category.Monoidal.Dialogue (Dialogue (..))
import Proarrow.Category.Monoidal.Hypergraph (cap)
import Proarrow.Core (CAT, CategoryOf (..), Is, Profunctor (..), Promonad (..), UN, dimapDefault, obj, type (+->))
import Proarrow.Functor (FunctorForRep (..))
import Proarrow.Limit.BinaryProduct (HasBinaryProducts (..))
import Proarrow.Limit.Terminal (HasTerminalObject (..))
import Proarrow.Monoid (CocommutativeComonoid, CommutativeMonoid, Comonoid (..), Monoid (..))
import Proarrow.Object (KnownListOf (..), withKnownAppend, withKnownAssoc, type (++))
import Proarrow.Optic (OPTIC (..), Prostrong)
import Proarrow.Optic.Lens (Lens, LensFl, lens)

-- | Objects are lists of dimensions, the product of the spaces of those dimensions.
type data DIFF = DF [Nat]

-- | A point of a product of spaces, or a change of one: a vector of coordinates for each space.
-- Taking a point of a product apart and putting it together handles the vectors without copying
-- their coordinates.
type Point :: [Nat] -> Type
newtype Point ns = Point {forall (ns :: [Nat]). Point ns -> [Vector Double]
vectors :: [SV.Vector Double]}

-- | A point of a single space.
pattern Vector :: SV.Vector Double -> Point '[n]
pattern $bVector :: forall (n :: Nat). Vector Double -> Point '[n]
$mVector :: forall {r} {n :: Nat}.
Point '[n] -> (Vector Double -> r) -> ((# #) -> r) -> r
Vector v = Point [v]

{-# COMPLETE Vector #-}

-- | The dimension of a space.
dim :: forall n. (KnownNat n) => P.Int
dim :: forall (n :: Nat). KnownNat n => Int
dim = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
P.fromIntegral (Proxy n -> Nat
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Nat
natVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @n))

-- | The coordinates of a point, one space after the other.
coordinates :: Point ns -> SV.Vector Double
coordinates :: forall (ns :: [Nat]). Point ns -> Vector Double
coordinates (Point [Vector Double]
vs) = [Vector Double] -> Vector Double
forall a. Storable a => [Vector a] -> Vector a
SV.concat [Vector Double]
vs

-- | The point with the given coordinates, one space after the other.
fromCoordinates :: forall ns. (KnownListOf KnownNat ns) => SV.Vector Double -> Point ns
fromCoordinates :: forall (ns :: [Nat]).
KnownListOf KnownNat ns =>
Vector Double -> Point ns
fromCoordinates = [Vector Double] -> Point ns
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point ([Vector Double] -> Point ns)
-> (Vector Double -> [Vector Double]) -> Vector Double -> Point ns
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. [Int] -> Vector Double -> [Vector Double]
forall {a}. Storable a => [Int] -> Vector a -> [Vector a]
go (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Int]
dimsOf @ns)
  where
    go :: [Int] -> Vector a -> [Vector a]
go [] Vector a
_ = []
    go (Int
d : [Int]
ds) Vector a
v = let (Vector a
l, Vector a
r) = Int -> Vector a -> (Vector a, Vector a)
forall a. Storable a => Int -> Vector a -> (Vector a, Vector a)
SV.splitAt Int
d Vector a
v in Vector a
l Vector a -> [Vector a] -> [Vector a]
forall a. a -> [a] -> [a]
: [Int] -> Vector a -> [Vector a]
go [Int]
ds Vector a
r

-- | The origin, which is also the change that changes nothing.
{-# INLINE zero #-}
zero :: forall ns. (KnownListOf KnownNat ns) => Point ns
zero :: forall (ns :: [Nat]). KnownListOf KnownNat ns => Point ns
zero = [Vector Double] -> Point ns
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point ((Int -> Vector Double) -> [Int] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
P.map (Int -> Double -> Vector Double
forall a. Storable a => Int -> a -> Vector a
`SV.replicate` Double
0) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Int]
dimsOf @ns))

-- | The sum of two points, or two changes.
{-# INLINE add #-}
add :: Point ns -> Point ns -> Point ns
add :: forall (ns :: [Nat]). Point ns -> Point ns -> Point ns
add (Point [Vector Double]
x) (Point [Vector Double]
y) = [Vector Double] -> Point ns
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point ((Vector Double -> Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
P.zipWith ((Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+)) [Vector Double]
x [Vector Double]
y)

-- | A point, or a change, multiplied by a number.
{-# INLINE scale #-}
scale :: Double -> Point ns -> Point ns
scale :: forall (ns :: [Nat]). Double -> Point ns -> Point ns
scale Double
s (Point [Vector Double]
x) = [Vector Double] -> Point ns
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point ((Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
P.map ((Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map (Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
*)) [Vector Double]
x)

-- | A point of a product as its two parts.
{-# INLINE split #-}
split :: forall as bs. (KnownListOf KnownNat as) => Point (as ++ bs) -> (Point as, Point bs)
split :: forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split = \(Point [Vector Double]
vs) -> let ([Vector Double]
l, [Vector Double]
r) = Int -> [Vector Double] -> ([Vector Double], [Vector Double])
forall a. Int -> [a] -> ([a], [a])
P.splitAt Int
n [Vector Double]
vs in ([Vector Double] -> Point as
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point [Vector Double]
l, [Vector Double] -> Point bs
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point [Vector Double]
r)
  where
    n :: Int
n = [Int] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
P.length (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Int]
dimsOf @as)

-- | The point of a product with the given parts.
{-# INLINE append #-}
append :: Point as -> Point bs -> Point (as ++ bs)
append :: forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append (Point [Vector Double]
l) (Point [Vector Double]
r) = [Vector Double] -> Point (as ++ bs)
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point ([Vector Double]
l [Vector Double] -> [Vector Double] -> [Vector Double]
forall a. [a] -> [a] -> [a]
P.++ [Vector Double]
r)

-- | A differentiable map: at a point, its image and its reverse derivative there, which takes a
-- change of the image to the change of the point it comes from.
type Diff :: CAT DIFF
data Diff a b where
  Diff
    :: forall as bs
     . (KnownListOf KnownNat as, KnownListOf KnownNat bs)
    => {forall (as :: [Nat]) (bs :: [Nat]).
Diff (DF as) (DF bs)
-> Point as -> (Point bs, Point bs -> Point as)
run :: Point as -> (Point bs, Point bs -> Point as)}
    -> Diff (DF as) (DF bs)

-- | The map of an arrow.
{-# INLINE forward #-}
forward :: DF as ~> DF bs -> Point as -> Point bs
forward :: forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs) -> Point as -> Point bs
forward DF as ~> DF bs
f Point as
x = (Point bs, Point bs -> Point as) -> Point bs
forall a b. (a, b) -> a
P.fst (Diff (DF as) (DF bs)
-> Point as -> (Point bs, Point bs -> Point as)
forall (as :: [Nat]) (bs :: [Nat]).
Diff (DF as) (DF bs)
-> Point as -> (Point bs, Point bs -> Point as)
run DF as ~> DF bs
Diff (DF as) (DF bs)
f Point as
x)

-- | The reverse derivative of an arrow at a point.
{-# INLINE backward #-}
backward :: DF as ~> DF bs -> Point as -> Point bs -> Point as
backward :: forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs) -> Point as -> Point bs -> Point as
backward DF as ~> DF bs
f Point as
x = (Point bs, Point bs -> Point as) -> Point bs -> Point as
forall a b. (a, b) -> b
P.snd (Diff (DF as) (DF bs)
-> Point as -> (Point bs, Point bs -> Point as)
forall (as :: [Nat]) (bs :: [Nat]).
Diff (DF as) (DF bs)
-> Point as -> (Point bs, Point bs -> Point as)
run DF as ~> DF bs
Diff (DF as) (DF bs)
f Point as
x)

-- | The arrow with the given map and reverse derivative, for when the two share no work.
{-# INLINE diff #-}
diff
  :: (KnownListOf KnownNat as, KnownListOf KnownNat bs)
  => (Point as -> Point bs) -> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff :: forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff Point as -> Point bs
f Point as -> Point bs -> Point as
fb = (Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff (\Point as
x -> (Point as -> Point bs
f Point as
x, Point as -> Point bs -> Point as
fb Point as
x))

-- | The identity between two lists of spaces that are the same list.
same :: forall as bs. (KnownListOf KnownNat as, KnownListOf KnownNat bs) => DF as ~> DF bs
same :: forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
DF as ~> DF bs
same = (Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff (\Point as
x -> (Point as -> Point bs
forall a b. Coercible a b => a -> b
coerce Point as
x, Point bs -> Point as
forall a b. Coercible a b => a -> b
coerce))

instance Profunctor Diff where
  dimap :: forall (c :: DIFF) (a :: DIFF) (b :: DIFF) (d :: DIFF).
(c ~> a) -> (b ~> d) -> Diff a b -> Diff c d
dimap = (c ~> a) -> (b ~> d) -> Diff a b -> Diff c d
Diff c a -> Diff b d -> Diff a b -> Diff c d
forall {k} (p :: CAT k) (c :: k) (a :: k) (b :: k) (d :: k).
Promonad p =>
p c a -> p b d -> p a b -> p c d
dimapDefault
  (Ob a, Ob b) => r
r \\ :: forall (a :: DIFF) (b :: DIFF) r.
((Ob a, Ob b) => r) -> Diff a b -> r
\\ Diff{} = r
(Ob a, Ob b) => r
r

instance Promonad Diff where
  id :: forall (a :: DIFF). Ob a => Diff a a
id = DF (UN DF a) ~> DF (UN DF a)
Diff a a
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
DF as ~> DF bs
same
  Diff Point as -> (Point bs, Point bs -> Point as)
g . :: forall (b :: DIFF) (c :: DIFF) (a :: DIFF).
Diff b c -> Diff a b -> Diff a c
. Diff Point as -> (Point bs, Point bs -> Point as)
f = (Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point as
x -> let (Point bs
y, Point bs -> Point as
fb) = Point as -> (Point bs, Point bs -> Point as)
f Point as
x; (Point bs
z, Point bs -> Point as
gb) = Point as -> (Point bs, Point bs -> Point as)
g Point as
Point bs
y in (Point bs
z, Point bs -> Point as
fb (Point bs -> Point as)
-> (Point bs -> Point bs) -> Point bs -> Point as
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Point bs -> Point as
Point bs -> Point bs
gb)

-- | Differentiable maps between products of real vector spaces.
instance CategoryOf DIFF where
  type (~>) = Diff
  type Ob a = (Is DF a, KnownListOf KnownNat (UN DF a))

instance MonoidalProfunctor Diff where
  one :: Diff Unit Unit
one = Diff Unit Unit
Diff (DF '[]) (DF '[])
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: DIFF). Ob a => Diff a a
id
  Diff @as @bs Point as -> (Point bs, Point bs -> Point as)
f ** :: forall (x1 :: DIFF) (x2 :: DIFF) (y1 :: DIFF) (y2 :: DIFF).
Diff x1 x2 -> Diff y1 y2 -> Diff (x1 ** y1) (x2 ** y2)
** Diff @cs @ds Point as -> (Point bs, Point bs -> Point as)
g =
    forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @cs ((KnownListOf KnownNat (as ++ as) => Diff (x1 ** y1) (x2 ** y2))
 -> Diff (x1 ** y1) (x2 ** y2))
-> (KnownListOf KnownNat (as ++ as) => Diff (x1 ** y1) (x2 ** y2))
-> Diff (x1 ** y1) (x2 ** y2)
forall a b. (a -> b) -> a -> b
$
      forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @bs @ds ((KnownListOf KnownNat (bs ++ bs) => Diff (x1 ** y1) (x2 ** y2))
 -> Diff (x1 ** y1) (x2 ** y2))
-> (KnownListOf KnownNat (bs ++ bs) => Diff (x1 ** y1) (x2 ** y2))
-> Diff (x1 ** y1) (x2 ** y2)
forall a b. (a -> b) -> a -> b
$
        (Point (as ++ as)
 -> (Point (bs ++ bs), Point (bs ++ bs) -> Point (as ++ as)))
-> Diff (DF (as ++ as)) (DF (bs ++ bs))
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point (as ++ as)
x ->
          let (Point as
a, Point as
c) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @as @cs Point (as ++ as)
x; (Point bs
b, Point bs -> Point as
fb) = Point as -> (Point bs, Point bs -> Point as)
f Point as
a; (Point bs
d, Point bs -> Point as
gb) = Point as -> (Point bs, Point bs -> Point as)
g Point as
c
          in (Point bs -> Point bs -> Point (bs ++ bs)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append Point bs
b Point bs
d, \Point (bs ++ bs)
dy -> let (Point bs
db, Point bs
dd) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @bs @ds Point (bs ++ bs)
dy in Point as -> Point as -> Point (as ++ as)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append (Point bs -> Point as
fb Point bs
db) (Point bs -> Point as
gb Point bs
dd))

-- | The product of spaces as the tensor.
instance Monoidal DIFF where
  type Unit = DF '[]
  type a ** b = DF (UN DF a ++ UN DF b)
  withOb2 :: forall (a :: DIFF) (b :: DIFF) r.
(Ob a, Ob b) =>
(Ob (a ** b) => r) -> r
withOb2 @a @b Ob (a ** b) => r
r = forall k (a :: k) (b :: k) r.
(HasBinaryProducts k, Ob a, Ob b) =>
(Ob (a && b) => r) -> r
withObProd @DIFF @a @b r
Ob (a ** b) => r
Ob (a && b) => r
r
  leftUnitor :: forall (a :: DIFF). Ob a => (Unit ** a) ~> a
leftUnitor = (Unit ** a) ~> a
Diff (DF (UN DF a)) (DF (UN DF a))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: DIFF). Ob a => Diff a a
id
  leftUnitorInv :: forall (a :: DIFF). Ob a => a ~> (Unit ** a)
leftUnitorInv = a ~> (Unit ** a)
Diff (DF (UN DF a)) (DF (UN DF a))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: DIFF). Ob a => Diff a a
id
  rightUnitor :: forall (a :: DIFF). Ob a => (a ** Unit) ~> a
rightUnitor @(DF as) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @'[] DF (UN DF a ++ '[]) ~> DF (UN DF a)
Diff (DF (UN DF a ++ '[])) (DF (UN DF a))
KnownListOf KnownNat (UN DF a ++ '[]) =>
Diff (DF (UN DF a ++ '[])) (DF (UN DF a))
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
DF as ~> DF bs
same
  rightUnitorInv :: forall (a :: DIFF). Ob a => a ~> (a ** Unit)
rightUnitorInv @(DF as) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @'[] DF (UN DF a) ~> DF (UN DF a ++ '[])
Diff (DF (UN DF a)) (DF (UN DF a ++ '[]))
KnownListOf KnownNat (UN DF a ++ '[]) =>
Diff (DF (UN DF a)) (DF (UN DF a ++ '[]))
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
DF as ~> DF bs
same
  associator :: forall (a :: DIFF) (b :: DIFF) (c :: DIFF).
(Ob a, Ob b, Ob c) =>
((a ** b) ** c) ~> (a ** (b ** c))
associator @(DF as) @(DF bs) @(DF cs) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) (cs :: [k]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) (cs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
withKnownAssoc @KnownNat @as @bs @cs DF ((UN DF a ++ UN DF b) ++ UN DF c)
~> DF (UN DF a ++ (UN DF b ++ UN DF c))
Diff
  (DF ((UN DF a ++ UN DF b) ++ UN DF c))
  (DF (UN DF a ++ (UN DF b ++ UN DF c)))
(KnownListOf KnownNat ((UN DF a ++ UN DF b) ++ UN DF c),
 KnownListOf KnownNat (UN DF a ++ (UN DF b ++ UN DF c))) =>
Diff
  (DF ((UN DF a ++ UN DF b) ++ UN DF c))
  (DF (UN DF a ++ (UN DF b ++ UN DF c)))
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
DF as ~> DF bs
same
  associatorInv :: forall (a :: DIFF) (b :: DIFF) (c :: DIFF).
(Ob a, Ob b, Ob c) =>
(a ** (b ** c)) ~> ((a ** b) ** c)
associatorInv @(DF as) @(DF bs) @(DF cs) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) (cs :: [k]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) (cs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
withKnownAssoc @KnownNat @as @bs @cs DF (UN DF a ++ (UN DF b ++ UN DF c))
~> DF ((UN DF a ++ UN DF b) ++ UN DF c)
Diff
  (DF (UN DF a ++ (UN DF b ++ UN DF c)))
  (DF ((UN DF a ++ UN DF b) ++ UN DF c))
(KnownListOf KnownNat ((UN DF a ++ UN DF b) ++ UN DF c),
 KnownListOf KnownNat (UN DF a ++ (UN DF b ++ UN DF c))) =>
Diff
  (DF (UN DF a ++ (UN DF b ++ UN DF c)))
  (DF ((UN DF a ++ UN DF b) ++ UN DF c))
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
DF as ~> DF bs
same

instance SymMonoidal DIFF where
  swap :: forall (a :: DIFF) (b :: DIFF).
(Ob a, Ob b) =>
(a ** b) ~> (b ** a)
swap @(DF as) @(DF bs) =
    forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @bs ((KnownListOf KnownNat (UN DF a ++ UN DF b) =>
  (a ** b) ~> (b ** a))
 -> (a ** b) ~> (b ** a))
-> (KnownListOf KnownNat (UN DF a ++ UN DF b) =>
    (a ** b) ~> (b ** a))
-> (a ** b) ~> (b ** a)
forall a b. (a -> b) -> a -> b
$
      forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @bs @as ((KnownListOf KnownNat (UN DF b ++ UN DF a) =>
  (a ** b) ~> (b ** a))
 -> (a ** b) ~> (b ** a))
-> (KnownListOf KnownNat (UN DF b ++ UN DF a) =>
    (a ** b) ~> (b ** a))
-> (a ** b) ~> (b ** a)
forall a b. (a -> b) -> a -> b
$
        (Point (UN DF a ++ UN DF b) -> Point (UN DF b ++ UN DF a))
-> (Point (UN DF a ++ UN DF b)
    -> Point (UN DF b ++ UN DF a) -> Point (UN DF a ++ UN DF b))
-> DF (UN DF a ++ UN DF b) ~> DF (UN DF b ++ UN DF a)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff
          (\Point (UN DF a ++ UN DF b)
x -> let (Point (UN DF a)
a, Point (UN DF b)
b) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @as @bs Point (UN DF a ++ UN DF b)
x in Point (UN DF b) -> Point (UN DF a) -> Point (UN DF b ++ UN DF a)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append Point (UN DF b)
b Point (UN DF a)
a)
          (\Point (UN DF a ++ UN DF b)
_ Point (UN DF b ++ UN DF a)
d -> let (Point (UN DF b)
db, Point (UN DF a)
da) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @bs @as Point (UN DF b ++ UN DF a)
d in Point (UN DF a) -> Point (UN DF b) -> Point (UN DF a ++ UN DF b)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append Point (UN DF a)
da Point (UN DF b)
db)

-- | The one-point space.
instance HasTerminalObject DIFF where
  type TerminalObject = DF '[]
  terminate :: forall (a :: DIFF). Ob a => a ~> TerminalObject
terminate = Point '[] -> DF (UN DF a) ~> DF '[]
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Point bs -> DF as ~> DF bs
constant Point '[]
forall (ns :: [Nat]). KnownListOf KnownNat ns => Point ns
zero

-- | The product of spaces: the reverse derivative of a pairing adds the changes of both parts.
instance HasBinaryProducts DIFF where
  type a && b = DF (UN DF a ++ UN DF b)
  withObProd :: forall (a :: DIFF) (b :: DIFF) r.
(Ob a, Ob b) =>
(Ob (a && b) => r) -> r
withObProd @(DF as) @(DF bs) Ob (a && b) => r
r = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @bs r
Ob (a && b) => r
KnownListOf KnownNat (UN DF a ++ UN DF b) => r
r
  fst :: forall (a :: DIFF) (b :: DIFF). (Ob a, Ob b) => (a && b) ~> a
fst @(DF as) @(DF bs) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @bs ((KnownListOf KnownNat (UN DF a ++ UN DF b) => (a && b) ~> a)
 -> (a && b) ~> a)
-> (KnownListOf KnownNat (UN DF a ++ UN DF b) => (a && b) ~> a)
-> (a && b) ~> a
forall a b. (a -> b) -> a -> b
$ (Point (UN DF a ++ UN DF b) -> Point (UN DF a))
-> (Point (UN DF a ++ UN DF b)
    -> Point (UN DF a) -> Point (UN DF a ++ UN DF b))
-> DF (UN DF a ++ UN DF b) ~> DF (UN DF a)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff ((Point (UN DF a), Point (UN DF b)) -> Point (UN DF a)
forall a b. (a, b) -> a
P.fst ((Point (UN DF a), Point (UN DF b)) -> Point (UN DF a))
-> (Point (UN DF a ++ UN DF b)
    -> (Point (UN DF a), Point (UN DF b)))
-> Point (UN DF a ++ UN DF b)
-> Point (UN DF a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @as @bs) (\Point (UN DF a ++ UN DF b)
_ Point (UN DF a)
da -> Point (UN DF a) -> Point (UN DF b) -> Point (UN DF a ++ UN DF b)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append Point (UN DF a)
da (forall (ns :: [Nat]). KnownListOf KnownNat ns => Point ns
zero @bs))
  snd :: forall (a :: DIFF) (b :: DIFF). (Ob a, Ob b) => (a && b) ~> b
snd @(DF as) @(DF bs) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @bs ((KnownListOf KnownNat (UN DF a ++ UN DF b) => (a && b) ~> b)
 -> (a && b) ~> b)
-> (KnownListOf KnownNat (UN DF a ++ UN DF b) => (a && b) ~> b)
-> (a && b) ~> b
forall a b. (a -> b) -> a -> b
$ (Point (UN DF a ++ UN DF b) -> Point (UN DF b))
-> (Point (UN DF a ++ UN DF b)
    -> Point (UN DF b) -> Point (UN DF a ++ UN DF b))
-> DF (UN DF a ++ UN DF b) ~> DF (UN DF b)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff ((Point (UN DF a), Point (UN DF b)) -> Point (UN DF b)
forall a b. (a, b) -> b
P.snd ((Point (UN DF a), Point (UN DF b)) -> Point (UN DF b))
-> (Point (UN DF a ++ UN DF b)
    -> (Point (UN DF a), Point (UN DF b)))
-> Point (UN DF a ++ UN DF b)
-> Point (UN DF b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @as @bs) (\Point (UN DF a ++ UN DF b)
_ Point (UN DF b)
db -> Point (UN DF a) -> Point (UN DF b) -> Point (UN DF a ++ UN DF b)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append (forall (ns :: [Nat]). KnownListOf KnownNat ns => Point ns
zero @as) Point (UN DF b)
db)
  Diff @_ @bs Point as -> (Point bs, Point bs -> Point as)
f &&& :: forall (a :: DIFF) (x :: DIFF) (y :: DIFF).
(a ~> x) -> (a ~> y) -> a ~> (x && y)
&&& Diff @_ @cs Point as -> (Point bs, Point bs -> Point as)
g =
    forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @bs @cs ((KnownListOf KnownNat (bs ++ bs) => a ~> (x && y))
 -> a ~> (x && y))
-> (KnownListOf KnownNat (bs ++ bs) => a ~> (x && y))
-> a ~> (x && y)
forall a b. (a -> b) -> a -> b
$
      (Point as -> (Point (bs ++ bs), Point (bs ++ bs) -> Point as))
-> Diff (DF as) (DF (bs ++ bs))
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point as
x ->
        let (Point bs
b, Point bs -> Point as
fb) = Point as -> (Point bs, Point bs -> Point as)
f Point as
Point as
x; (Point bs
c, Point bs -> Point as
gb) = Point as -> (Point bs, Point bs -> Point as)
g Point as
x
        in (Point bs -> Point bs -> Point (bs ++ bs)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append Point bs
b Point bs
c, \Point (bs ++ bs)
d -> let (Point bs
db, Point bs
dc) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @bs @cs Point (bs ++ bs)
d in Point as -> Point as -> Point as
forall (ns :: [Nat]). Point ns -> Point ns -> Point ns
add (Point bs -> Point as
fb Point bs
db) (Point bs -> Point as
gb Point bs
dc))

-- | Copying a point: the changes that come back are added.
instance (KnownListOf KnownNat ns) => Comonoid (DF ns :: DIFF) where
  counit :: DF ns ~> Unit
counit = DF ns ~> Unit
DF ns ~> TerminalObject
forall k (a :: k).
(HasTerminalObject k, Ob a) =>
a ~> TerminalObject
forall (a :: DIFF). Ob a => a ~> TerminalObject
terminate
  comult :: DF ns ~> (DF ns ** DF ns)
comult = DF ns ~> DF ns
Diff (DF ns) (DF ns)
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: DIFF). Ob a => Diff a a
id (DF ns ~> DF ns) -> (DF ns ~> DF ns) -> DF ns ~> (DF ns && DF ns)
forall k (a :: k) (x :: k) (y :: k).
HasBinaryProducts k =>
(a ~> x) -> (a ~> y) -> a ~> (x && y)
forall (a :: DIFF) (x :: DIFF) (y :: DIFF).
(a ~> x) -> (a ~> y) -> a ~> (x && y)
&&& DF ns ~> DF ns
Diff (DF ns) (DF ns)
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: DIFF). Ob a => Diff a a
id

instance (KnownListOf KnownNat ns) => CocommutativeComonoid (DF ns :: DIFF)
instance CopyDiscard DIFF

-- | Adding vectors: the change of a sum is passed back to both terms.
instance (KnownListOf KnownNat ns) => Monoid (DF ns :: DIFF) where
  mempty :: Unit ~> DF ns
mempty = Point ns -> DF '[] ~> DF ns
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Point bs -> DF as ~> DF bs
constant Point ns
forall (ns :: [Nat]). KnownListOf KnownNat ns => Point ns
zero
  mappend :: (DF ns ** DF ns) ~> DF ns
mappend = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @ns @ns ((KnownListOf KnownNat (ns ++ ns) => (DF ns ** DF ns) ~> DF ns)
 -> (DF ns ** DF ns) ~> DF ns)
-> (KnownListOf KnownNat (ns ++ ns) => (DF ns ** DF ns) ~> DF ns)
-> (DF ns ** DF ns) ~> DF ns
forall a b. (a -> b) -> a -> b
$ (Point (ns ++ ns) -> Point ns)
-> (Point (ns ++ ns) -> Point ns -> Point (ns ++ ns))
-> DF (ns ++ ns) ~> DF ns
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff (\Point (ns ++ ns)
x -> let (Point ns
a, Point ns
b) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @ns @ns Point (ns ++ ns)
x in Point ns -> Point ns -> Point ns
forall (ns :: [Nat]). Point ns -> Point ns -> Point ns
add Point ns
a Point ns
b) (\Point (ns ++ ns)
_ Point ns
d -> Point ns -> Point ns -> Point (ns ++ ns)
forall (as :: [Nat]) (bs :: [Nat]).
Point as -> Point bs -> Point (as ++ bs)
append Point ns
d Point ns
d)

instance (KnownListOf KnownNat ns) => CommutativeMonoid (DF ns :: DIFF)

-- | A linear map, given as a tensor network, its input and result each the vector of a tensor's
-- entries. The change of the input is the change of the result through the 'dual' of the map.
linear
  :: forall as bs
   . (KnownListOf KnownNat as, KnownListOf KnownNat bs, KnownNat (Size as), KnownNat (Size bs))
  => (TN as :: TNET Double) ~> TN bs
  -> DF '[Size as] ~> DF '[Size bs]
linear :: forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs,
 KnownNat (Size as), KnownNat (Size bs)) =>
(TN as ~> TN bs) -> DF '[Size as] ~> DF '[Size bs]
linear TN as ~> TN bs
f =
  (Point '[Size as] -> Point '[Size bs])
-> (Point '[Size as] -> Point '[Size bs] -> Point '[Size as])
-> DF '[Size as] ~> DF '[Size bs]
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff
    (\(Vector Vector Double
x) -> Vector Double -> Point '[Size bs]
forall (n :: Nat). Vector Double -> Point '[n]
Vector (TensorNetwork (TN '[]) (TN bs) -> Vector Double
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector (TN as ~> TN bs
TensorNetwork (TN as) (TN bs)
f TensorNetwork (TN as) (TN bs)
-> TensorNetwork (TN '[]) (TN as) -> TensorNetwork (TN '[]) (TN bs)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. forall (as :: [Nat]) (bs :: [Nat]).
(Scalar Double, KnownListOf KnownNat as,
 KnownListOf KnownNat bs) =>
Vector Double -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector @'[] Vector Double
x)))
    (\Point '[Size as]
_ (Vector Vector Double
dz) -> Vector Double -> Point '[Size as]
forall (n :: Nat). Vector Double -> Point '[n]
Vector (TensorNetwork (TN '[]) (TN as) -> Vector Double
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector ((TN as ~> TN bs) -> Dual (TN bs) ~> Dual (TN as)
forall k (a :: k) (b :: k).
Dialogue k =>
(a ~> b) -> Dual b ~> Dual a
forall (a :: TNET Double) (b :: TNET Double).
(a ~> b) -> Dual b ~> Dual a
dual TN as ~> TN bs
f TensorNetwork (TN bs) (TN as)
-> TensorNetwork (TN '[]) (TN bs) -> TensorNetwork (TN '[]) (TN as)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. forall (as :: [Nat]) (bs :: [Nat]).
(Scalar Double, KnownListOf KnownNat as,
 KnownListOf KnownNat bs) =>
Vector Double -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector @'[] Vector Double
dz)))

-- | A function applied to each coordinate, given with its derivative.
pointwise :: (KnownNat n) => (Double -> Double) -> (Double -> Double) -> DF '[n] ~> DF '[n]
pointwise :: forall (n :: Nat).
KnownNat n =>
(Double -> Double) -> (Double -> Double) -> DF '[n] ~> DF '[n]
pointwise Double -> Double
f Double -> Double
f' = (Point '[n] -> Point '[n])
-> (Point '[n] -> Point '[n] -> Point '[n]) -> DF '[n] ~> DF '[n]
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff (\(Vector Vector Double
x) -> Vector Double -> Point '[n]
forall (n :: Nat). Vector Double -> Point '[n]
Vector ((Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map Double -> Double
f Vector Double
x)) (\(Vector Vector Double
x) (Vector Vector Double
d) -> Vector Double -> Point '[n]
forall (n :: Nat). Vector Double -> Point '[n]
Vector ((Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) ((Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map Double -> Double
f' Vector Double
x) Vector Double
d))

-- | The coordinatewise product of two vectors: 'bilinear' with the merging spider, written out.
times :: forall n. (KnownNat n) => DF '[n, n] ~> DF '[n]
times :: forall (n :: Nat). KnownNat n => DF '[n, n] ~> DF '[n]
times = (Point '[n, n] -> (Point '[n], Point '[n] -> Point '[n, n]))
-> Diff (DF '[n, n]) (DF '[n])
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point '[n, n]
p ->
  let (Vector Vector Double
x, Vector Vector Double
y) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @'[n] @'[n] Point '[n, n]
Point ('[n] ++ '[n])
p
  in (Vector Double -> Point '[n]
forall (n :: Nat). Vector Double -> Point '[n]
Vector ((Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
x Vector Double
y), \(Vector Vector Double
d) -> [Vector Double] -> Point '[n, n]
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point [(Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
d Vector Double
y, (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
d Vector Double
x])

-- | The dot product of two vectors: 'bilinear' with the cap, which sums the products.
dot :: forall n. (KnownNat n) => DF '[n, n] ~> DF '[1]
dot :: forall (n :: Nat). KnownNat n => DF '[n, n] ~> DF '[1]
dot = forall (as :: [Nat]) (bs :: [Nat]) (cs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs,
 KnownListOf KnownNat cs, KnownNat (Size as), KnownNat (Size bs),
 KnownNat (Size cs)) =>
(TN (as ++ bs) ~> TN cs) -> DF '[Size as, Size bs] ~> DF '[Size cs]
bilinear @'[n] @'[n] (forall {k} (a :: k). Frobenius a => (a ** a) ~> Unit
forall (a :: TNET Double). Frobenius a => (a ** a) ~> Unit
cap @(TN '[n]))

-- | The product of an @m@ by @n@ matrix, given row by row, with a vector: 'matMat' with a single
-- column.
matVec :: forall m n. (KnownNat m, KnownNat n, KnownNat (m * n)) => DF '[m * n, n] ~> DF '[m]
matVec :: forall (m :: Nat) (n :: Nat).
(KnownNat m, KnownNat n, KnownNat (m * n)) =>
DF '[m * n, n] ~> DF '[m]
matVec = forall (m :: Nat) (n :: Nat) (k :: Nat).
(KnownNat m, KnownNat n, KnownNat k, KnownNat (m * n),
 KnownNat (n * k), KnownNat (m * k)) =>
DF '[m * n, n * k] ~> DF '[m * k]
matMat @m @n @1

-- | The product of an @m@ by @n@ matrix with an @n@ by @k@ matrix, both given row by row. With a
-- batch of @m@ vectors as the rows of the first, the second is a linear map applied to each of them.
-- It is 'bilinear' with the cap on the summed index, written out with 'matMul': the change of the
-- first matrix is the change of the result times the transposed second, and the change of the second
-- is the transposed first times the change of the result, which adds up the changes from the whole
-- batch.
matMat
  :: forall m n k
   . (KnownNat m, KnownNat n, KnownNat k, KnownNat (m * n), KnownNat (n * k), KnownNat (m * k))
  => DF '[m * n, n * k] ~> DF '[m * k]
matMat :: forall (m :: Nat) (n :: Nat) (k :: Nat).
(KnownNat m, KnownNat n, KnownNat k, KnownNat (m * n),
 KnownNat (n * k), KnownNat (m * k)) =>
DF '[m * n, n * k] ~> DF '[m * k]
matMat = (Point '[m * n, n * k]
 -> (Point '[m * k], Point '[m * k] -> Point '[m * n, n * k]))
-> Diff (DF '[m * n, n * k]) (DF '[m * k])
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point '[m * n, n * k]
p ->
  let (Vector Vector Double
x, Vector Vector Double
w) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @'[m * n] @'[n * k] Point '[m * n, n * k]
Point ('[m * n] ++ '[n * k])
p
  in ( Vector Double -> Point '[m * k]
forall (n :: Nat). Vector Double -> Point '[n]
Vector (Bool
-> Bool
-> Int
-> Int
-> Int
-> Vector Double
-> Vector Double
-> Vector Double
forall e.
Scalar e =>
Bool
-> Bool -> Int -> Int -> Int -> Vector e -> Vector e -> Vector e
matMul Bool
P.False Bool
P.False Int
m Int
k Int
n Vector Double
x Vector Double
w)
     , \(Vector Vector Double
dy) -> [Vector Double] -> Point '[m * n, n * k]
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point [Bool
-> Bool
-> Int
-> Int
-> Int
-> Vector Double
-> Vector Double
-> Vector Double
forall e.
Scalar e =>
Bool
-> Bool -> Int -> Int -> Int -> Vector e -> Vector e -> Vector e
matMul Bool
P.False Bool
P.True Int
m Int
n Int
k Vector Double
dy Vector Double
w, Bool
-> Bool
-> Int
-> Int
-> Int
-> Vector Double
-> Vector Double
-> Vector Double
forall e.
Scalar e =>
Bool
-> Bool -> Int -> Int -> Int -> Vector e -> Vector e -> Vector e
matMul Bool
P.True Bool
P.False Int
n Int
k Int
m Vector Double
x Vector Double
dy]
     )
  where
    m :: Int
m = forall (n :: Nat). KnownNat n => Int
dim @m
    n :: Int
n = forall (n :: Nat). KnownNat n => Int
dim @n
    k :: Int
k = forall (n :: Nat). KnownNat n => Int
dim @k

-- | A vector copied to each of the @b@ rows of a matrix, given row by row: for adding the same
-- vector to each of a batch of @b@ vectors. The change of the vector adds up the rows of the change
-- of the matrix.
broadcast :: forall m b. (KnownNat m, KnownNat b, KnownNat (b * m)) => DF '[m] ~> DF '[b * m]
broadcast :: forall (m :: Nat) (b :: Nat).
(KnownNat m, KnownNat b, KnownNat (b * m)) =>
DF '[m] ~> DF '[b * m]
broadcast =
  (Point '[m] -> Point '[b * m])
-> (Point '[m] -> Point '[b * m] -> Point '[m])
-> DF '[m] ~> DF '[b * m]
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> Point bs)
-> (Point as -> Point bs -> Point as) -> DF as ~> DF bs
diff
    (\(Vector Vector Double
v) -> Vector Double -> Point '[b * m]
forall (n :: Nat). Vector Double -> Point '[n]
Vector ([Vector Double] -> Vector Double
forall a. Storable a => [Vector a] -> Vector a
SV.concat (Int -> Vector Double -> [Vector Double]
forall a. Int -> a -> [a]
P.replicate Int
b Vector Double
v)))
    (\Point '[m]
_ (Vector Vector Double
d) -> Vector Double -> Point '[m]
forall (n :: Nat). Vector Double -> Point '[n]
Vector ((Vector Double -> Vector Double -> Vector Double)
-> Vector Double -> [Vector Double] -> Vector Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
P.foldl' ((Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+)) (Int -> Double -> Vector Double
forall a. Storable a => Int -> a -> Vector a
SV.replicate Int
m Double
0) (Int -> Int -> Vector Double -> [Vector Double]
rows Int
b Int
m Vector Double
d)))
  where
    m :: Int
m = forall (n :: Nat). KnownNat n => Int
dim @m
    b :: Int
b = forall (n :: Nat). KnownNat n => Int
dim @b

-- | The softmax of each of the @b@ rows of a matrix of scores, given row by row: the exponentials of
-- a row divided by their sum, for each of a batch of @b@ vectors a probability for each of @k@
-- classes. The change of a row of scores is the probabilities times the change of the row minus its
-- average under them.
softmax :: forall k b. (KnownNat k, KnownNat b, KnownNat (b * k)) => DF '[b * k] ~> DF '[b * k]
softmax :: forall (k :: Nat) (b :: Nat).
(KnownNat k, KnownNat b, KnownNat (b * k)) =>
DF '[b * k] ~> DF '[b * k]
softmax = (Point '[b * k]
 -> (Point '[b * k], Point '[b * k] -> Point '[b * k]))
-> Diff (DF '[b * k]) (DF '[b * k])
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \(Vector Vector Double
z) ->
  let ps :: [Vector Double]
ps = (Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
P.map ((Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map Double -> Double
forall a. Floating a => a -> a
P.exp (Vector Double -> Vector Double)
-> (Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Vector Double -> Vector Double
logSoftmax) (Int -> Int -> Vector Double -> [Vector Double]
rows Int
b Int
k Vector Double
z)
  in ( Vector Double -> Point '[b * k]
forall (n :: Nat). Vector Double -> Point '[n]
Vector ([Vector Double] -> Vector Double
forall a. Storable a => [Vector a] -> Vector a
SV.concat [Vector Double]
ps)
     , \(Vector Vector Double
d) ->
         Vector Double -> Point '[b * k]
forall (n :: Nat). Vector Double -> Point '[n]
Vector
           ( [Vector Double] -> Vector Double
forall a. Storable a => [Vector a] -> Vector a
SV.concat
               ( (Vector Double -> Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
P.zipWith
                   (\Vector Double
p Vector Double
dr -> let average :: Double
average = Vector Double -> Double
forall a. (Storable a, Num a) => Vector a -> a
SV.sum ((Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
dr Vector Double
p) in (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith (\Double
q Double
e -> Double
q Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
e Double -> Double -> Double
forall a. Num a => a -> a -> a
P.- Double
average)) Vector Double
p Vector Double
dr)
                   [Vector Double]
ps
                   (Int -> Int -> Vector Double -> [Vector Double]
rows Int
b Int
k Vector Double
d)
               )
           )
     )
  where
    k :: Int
k = forall (n :: Nat). KnownNat n => Int
dim @k
    b :: Int
b = forall (n :: Nat). KnownNat n => Int
dim @b

-- | The cross-entropy of the softmax of each of the @b@ rows of a matrix of scores against the same
-- row of a matrix of target probabilities, both given row by row, summed over the rows: the loss of
-- a classifier on a batch. With the softmax computed along, the change of the scores is the
-- probabilities minus the targets, for targets that sum to 1, and the change of the targets is
-- minus the logarithms of the probabilities.
softmaxCrossEntropy :: forall k b. (KnownNat k, KnownNat b, KnownNat (b * k)) => DF '[b * k, b * k] ~> DF '[1]
softmaxCrossEntropy :: forall (k :: Nat) (b :: Nat).
(KnownNat k, KnownNat b, KnownNat (b * k)) =>
DF '[b * k, b * k] ~> DF '[1]
softmaxCrossEntropy = (Point '[b * k, b * k]
 -> (Point '[1], Point '[1] -> Point '[b * k, b * k]))
-> Diff (DF '[b * k, b * k]) (DF '[1])
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point '[b * k, b * k]
x ->
  let (Vector Vector Double
z, Vector Vector Double
t) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @'[b * k] @'[b * k] Point '[b * k, b * k]
Point ('[b * k] ++ '[b * k])
x
      logPs :: [Vector Double]
logPs = (Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
P.map Vector Double -> Vector Double
logSoftmax (Int -> Int -> Vector Double -> [Vector Double]
rows Int
b Int
k Vector Double
z)
      logP :: Vector Double
logP = [Vector Double] -> Vector Double
forall a. Storable a => [Vector a] -> Vector a
SV.concat [Vector Double]
logPs
  in ( Vector Double -> Point '[1]
forall (n :: Nat). Vector Double -> Point '[n]
Vector (Double -> Vector Double
forall a. Storable a => a -> Vector a
SV.singleton (Double -> Double
forall a. Num a => a -> a
P.negate (Vector Double -> Double
forall a. (Storable a, Num a) => Vector a -> a
SV.sum ((Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
t Vector Double
logP))))
     , \(Vector Vector Double
d) ->
         let s :: Double
s = Vector Double -> Double
forall a. Storable a => Vector a -> a
SV.head Vector Double
d
             scores :: Vector Double -> Vector Double -> Vector Double
scores Vector Double
lp Vector Double
tr = let mass :: Double
mass = Vector Double -> Double
forall a. (Storable a, Num a) => Vector a -> a
SV.sum Vector Double
tr in (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
SV.zipWith (\Double
l Double
e -> Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double
forall a. Floating a => a -> a
P.exp Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mass Double -> Double -> Double
forall a. Num a => a -> a -> a
P.- Double
e)) Vector Double
lp Vector Double
tr
         in [Vector Double] -> Point '[b * k, b * k]
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point [[Vector Double] -> Vector Double
forall a. Storable a => [Vector a] -> Vector a
SV.concat ((Vector Double -> Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
P.zipWith Vector Double -> Vector Double -> Vector Double
scores [Vector Double]
logPs (Int -> Int -> Vector Double -> [Vector Double]
rows Int
b Int
k Vector Double
t)), (Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map (\Double
l -> Double -> Double
forall a. Num a => a -> a
P.negate (Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
l)) Vector Double
logP]
     )
  where
    k :: Int
k = forall (n :: Nat). KnownNat n => Int
dim @k
    b :: Int
b = forall (n :: Nat). KnownNat n => Int
dim @b

-- | A bilinear map, given as a tensor network from the tensor product of its two inputs, each input
-- and the result the vector of a tensor's entries: 'linear' in each input. The change of one input
-- is the change of the result through the 'dual' of the map with the other input fixed.
bilinear
  :: forall as bs cs
   . ( KnownListOf KnownNat as
     , KnownListOf KnownNat bs
     , KnownListOf KnownNat cs
     , KnownNat (Size as)
     , KnownNat (Size bs)
     , KnownNat (Size cs)
     )
  => (TN (as ++ bs) :: TNET Double) ~> TN cs
  -> DF '[Size as, Size bs] ~> DF '[Size cs]
bilinear :: forall (as :: [Nat]) (bs :: [Nat]) (cs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs,
 KnownListOf KnownNat cs, KnownNat (Size as), KnownNat (Size bs),
 KnownNat (Size cs)) =>
(TN (as ++ bs) ~> TN cs) -> DF '[Size as, Size bs] ~> DF '[Size cs]
bilinear TN (as ++ bs) ~> TN cs
f = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @bs ((KnownListOf KnownNat (as ++ bs) =>
  DF '[Size as, Size bs] ~> DF '[Size cs])
 -> DF '[Size as, Size bs] ~> DF '[Size cs])
-> (KnownListOf KnownNat (as ++ bs) =>
    DF '[Size as, Size bs] ~> DF '[Size cs])
-> DF '[Size as, Size bs] ~> DF '[Size cs]
forall a b. (a -> b) -> a -> b
$ (Point '[Size as, Size bs]
 -> (Point '[Size cs],
     Point '[Size cs] -> Point '[Size as, Size bs]))
-> Diff (DF '[Size as, Size bs]) (DF '[Size cs])
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff \Point '[Size as, Size bs]
p ->
  let (Vector Vector Double
x, Vector Vector Double
y) = forall (as :: [Nat]) (bs :: [Nat]).
KnownListOf KnownNat as =>
Point (as ++ bs) -> (Point as, Point bs)
split @'[Size as] @'[Size bs] Point '[Size as, Size bs]
Point ('[Size as] ++ '[Size bs])
p
      tx :: TensorNetwork (TN '[]) (TN as)
tx = forall (as :: [Nat]) (bs :: [Nat]).
(Scalar Double, KnownListOf KnownNat as,
 KnownListOf KnownNat bs) =>
Vector Double -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector @'[] @as Vector Double
x
      ty :: TensorNetwork (TN '[]) (TN bs)
ty = forall (as :: [Nat]) (bs :: [Nat]).
(Scalar Double, KnownListOf KnownNat as,
 KnownListOf KnownNat bs) =>
Vector Double -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector @'[] @bs Vector Double
y
      withY :: TensorNetwork (TN as) (TN cs)
withY = TN (as ++ bs) ~> TN cs
TensorNetwork (TN (as ++ bs)) (TN cs)
f TensorNetwork (TN (as ++ bs)) (TN cs)
-> TensorNetwork (TN as) (TN (as ++ bs))
-> TensorNetwork (TN as) (TN cs)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. (forall {k} (a :: k). (CategoryOf k, Ob a) => Obj a
forall (a :: TNET Double).
(CategoryOf (TNET Double), Ob a) =>
Obj a
obj @(TN as) TensorNetwork (TN as) (TN as)
-> TensorNetwork (TN '[]) (TN bs)
-> TensorNetwork (TN as ** TN '[]) (TN as ** TN bs)
forall {j} {k} (p :: j +-> k) (x1 :: k) (x2 :: j) (y1 :: k)
       (y2 :: j).
MonoidalProfunctor p =>
p x1 x2 -> p y1 y2 -> p (x1 ** y1) (x2 ** y2)
forall (x1 :: TNET Double) (x2 :: TNET Double) (y1 :: TNET Double)
       (y2 :: TNET Double).
TensorNetwork x1 x2
-> TensorNetwork y1 y2 -> TensorNetwork (x1 ** y1) (x2 ** y2)
** TensorNetwork (TN '[]) (TN bs)
ty) TensorNetwork (TN (UN TN (TN as) ++ '[])) (TN (as ++ bs))
-> TensorNetwork (TN as) (TN (UN TN (TN as) ++ '[]))
-> TensorNetwork (TN as) (TN (as ++ bs))
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. forall k (a :: k). (Monoidal k, Ob a) => a ~> (a ** Unit)
rightUnitorInv @_ @(TN as)
      withX :: TensorNetwork (TN (UN TN (TN bs))) (TN cs)
withX = TN (as ++ bs) ~> TN cs
TensorNetwork (TN (as ++ bs)) (TN cs)
f TensorNetwork (TN (as ++ bs)) (TN cs)
-> TensorNetwork (TN (UN TN (TN bs))) (TN (as ++ bs))
-> TensorNetwork (TN (UN TN (TN bs))) (TN cs)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. (TensorNetwork (TN '[]) (TN as)
tx TensorNetwork (TN '[]) (TN as)
-> TensorNetwork (TN bs) (TN bs)
-> TensorNetwork (TN '[] ** TN bs) (TN as ** TN bs)
forall {j} {k} (p :: j +-> k) (x1 :: k) (x2 :: j) (y1 :: k)
       (y2 :: j).
MonoidalProfunctor p =>
p x1 x2 -> p y1 y2 -> p (x1 ** y1) (x2 ** y2)
forall (x1 :: TNET Double) (x2 :: TNET Double) (y1 :: TNET Double)
       (y2 :: TNET Double).
TensorNetwork x1 x2
-> TensorNetwork y1 y2 -> TensorNetwork (x1 ** y1) (x2 ** y2)
** forall {k} (a :: k). (CategoryOf k, Ob a) => Obj a
forall (a :: TNET Double).
(CategoryOf (TNET Double), Ob a) =>
Obj a
obj @(TN bs))
  in ( Vector Double -> Point '[Size cs]
forall (n :: Nat). Vector Double -> Point '[n]
Vector (TensorNetwork (TN '[]) (TN cs) -> Vector Double
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector (TN (as ++ bs) ~> TN cs
TensorNetwork (TN (as ++ bs)) (TN cs)
f TensorNetwork (TN (as ++ bs)) (TN cs)
-> TensorNetwork (TN '[]) (TN (as ++ bs))
-> TensorNetwork (TN '[]) (TN cs)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. (TensorNetwork (TN '[]) (TN as)
tx TensorNetwork (TN '[]) (TN as)
-> TensorNetwork (TN '[]) (TN bs)
-> TensorNetwork (TN '[] ** TN '[]) (TN as ** TN bs)
forall {j} {k} (p :: j +-> k) (x1 :: k) (x2 :: j) (y1 :: k)
       (y2 :: j).
MonoidalProfunctor p =>
p x1 x2 -> p y1 y2 -> p (x1 ** y1) (x2 ** y2)
forall (x1 :: TNET Double) (x2 :: TNET Double) (y1 :: TNET Double)
       (y2 :: TNET Double).
TensorNetwork x1 x2
-> TensorNetwork y1 y2 -> TensorNetwork (x1 ** y1) (x2 ** y2)
** TensorNetwork (TN '[]) (TN bs)
ty)))
     , \(Vector Vector Double
dz) ->
         let tdz :: TensorNetwork (TN '[]) (TN cs)
tdz = forall (as :: [Nat]) (bs :: [Nat]).
(Scalar Double, KnownListOf KnownNat as,
 KnownListOf KnownNat bs) =>
Vector Double -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector @'[] Vector Double
dz
         in [Vector Double] -> Point '[Size as, Size bs]
forall (ns :: [Nat]). [Vector Double] -> Point ns
Point [TensorNetwork (TN '[]) (TN as) -> Vector Double
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector ((TN as ~> TN cs) -> Dual (TN cs) ~> Dual (TN as)
forall k (a :: k) (b :: k).
Dialogue k =>
(a ~> b) -> Dual b ~> Dual a
forall (a :: TNET Double) (b :: TNET Double).
(a ~> b) -> Dual b ~> Dual a
dual TN as ~> TN cs
TensorNetwork (TN as) (TN cs)
withY TensorNetwork (TN cs) (TN as)
-> TensorNetwork (TN '[]) (TN cs) -> TensorNetwork (TN '[]) (TN as)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. TensorNetwork (TN '[]) (TN cs)
tdz), TensorNetwork (TN '[]) (TN bs) -> Vector Double
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector ((TN bs ~> TN cs) -> Dual (TN cs) ~> Dual (TN bs)
forall k (a :: k) (b :: k).
Dialogue k =>
(a ~> b) -> Dual b ~> Dual a
forall (a :: TNET Double) (b :: TNET Double).
(a ~> b) -> Dual b ~> Dual a
dual TN bs ~> TN cs
TensorNetwork (TN (UN TN (TN bs))) (TN cs)
withX TensorNetwork (TN cs) (TN bs)
-> TensorNetwork (TN '[]) (TN cs) -> TensorNetwork (TN '[]) (TN bs)
forall {k} (p :: CAT k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
forall (b :: TNET Double) (c :: TNET Double) (a :: TNET Double).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. TensorNetwork (TN '[]) (TN cs)
tdz)]
     )

-- | The logarithms of the softmax of a vector, with its largest entry taken out before the
-- exponentials so that they cannot overflow.
logSoftmax :: SV.Vector Double -> SV.Vector Double
logSoftmax :: Vector Double -> Vector Double
logSoftmax Vector Double
z = (Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map (Double -> Double -> Double
forall a. Num a => a -> a -> a
P.subtract Double
logSum) Vector Double
z
  where
    m :: Double
m = Vector Double -> Double
forall a. (Storable a, Ord a) => Vector a -> a
SV.maximum Vector Double
z
    logSum :: Double
logSum = Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
P.log (Vector Double -> Double
forall a. (Storable a, Num a) => Vector a -> a
SV.sum ((Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
SV.map (\Double
x -> Double -> Double
forall a. Floating a => a -> a
P.exp (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
P.- Double
m)) Vector Double
z))

-- | The @b@ rows of length @k@ of a matrix given row by row.
rows :: P.Int -> P.Int -> SV.Vector Double -> [SV.Vector Double]
rows :: Int -> Int -> Vector Double -> [Vector Double]
rows Int
b Int
k Vector Double
v = [Int -> Int -> Vector Double -> Vector Double
forall a. Storable a => Int -> Int -> Vector a -> Vector a
SV.slice (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int
k Vector Double
v | Int
j <- Int -> Int -> [Int]
forall a. Enum a => a -> a -> [a]
P.enumFromTo Int
0 (Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
P.- Int
1)]

-- | The map to a fixed point, whose reverse derivative is the zero change.
constant :: forall as bs. (KnownListOf KnownNat as, KnownListOf KnownNat bs) => Point bs -> DF as ~> DF bs
constant :: forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Point bs -> DF as ~> DF bs
constant Point bs
p = (Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
forall (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Point as -> (Point bs, Point bs -> Point as))
-> Diff (DF as) (DF bs)
Diff ((Point bs, Point bs -> Point as)
-> Point as -> (Point bs, Point bs -> Point as)
forall a b. a -> b -> a
P.const (Point bs
p, Point as -> Point bs -> Point as
forall a b. a -> b -> a
P.const Point as
forall (ns :: [Nat]). KnownListOf KnownNat ns => Point ns
zero))

-- | An arrow as a lens: the map is the view, and the reverse derivative the update. Composing
-- arrows and composing their lenses agree, which is the chain rule; 'AsLens' is the functor. A lens
-- views and updates separately, so the update runs the map again.
toLens :: DF as ~> DF bs -> Lens (Point as) (Point as) (Point bs) (Point bs)
toLens :: forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs)
-> Lens (Point as) (Point as) (Point bs) (Point bs)
toLens DF as ~> DF bs
f = (Point as ~> Point bs)
-> ((Point as && Point bs) ~> Point as)
-> Lens (Point as) (Point as) (Point bs) (Point bs)
forall {k} (s :: k) (t :: k) (a :: k) (b :: k).
(HasBinaryProducts k, Ob b) =>
(s ~> a) -> ((s && b) ~> t) -> Lens s t a b
lens ((DF as ~> DF bs) -> Point as -> Point bs
forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs) -> Point as -> Point bs
forward DF as ~> DF bs
f) ((Point as -> Point bs -> Point as)
-> (Point as, Point bs) -> Point as
forall a b c. (a -> b -> c) -> (a, b) -> c
P.uncurry ((DF as ~> DF bs) -> Point as -> Point bs -> Point as
forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs) -> Point as -> Point bs -> Point as
backward DF as ~> DF bs
f))

-- | The underlying functions: a differentiable map as a function from the points of its source to
-- those of its target, forgetting the reverse derivative.
data family Forget :: DIFF +-> Type

instance FunctorForRep Forget where
  type Forget @ a = Point (UN DF a)
  fmap :: forall (a :: DIFF) (b :: DIFF).
(a ~> b) -> (Forget @ a) ~> (Forget @ b)
fmap f :: a ~> b
f@Diff{} = (DF as ~> DF bs) -> Point as -> Point bs
forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs) -> Point as -> Point bs
forward a ~> b
DF as ~> DF bs
f

-- | Arrows as lenses, by 'toLens'. The lens of an arrow goes the other way, so this is a functor into
-- the opposite of the category of lenses, and that it is a functor is the chain rule.
data family AsLens :: DIFF +-> OPPOSITE (OPTIC Type Type (Prostrong LensFl))

instance FunctorForRep AsLens where
  type AsLens @ a = OP (OPT (Point (UN DF a)) (Point (UN DF a)))
  fmap :: forall (a :: DIFF) (b :: DIFF).
(a ~> b) -> (AsLens @ a) ~> (AsLens @ b)
fmap f :: a ~> b
f@Diff{} = Optic_ (OPT (Point bs) (Point bs)) (OPT (Point as) (Point as))
-> Op
     Optic_
     (OP (OPT (Point as) (Point as)))
     (OP (OPT (Point bs) (Point bs)))
forall {j} {k} (p :: j +-> k) (b1 :: k) (a1 :: j).
p b1 a1 -> Op p (OP a1) (OP b1)
Op ((DF as ~> DF bs)
-> Optic_ (OPT (Point bs) (Point bs)) (OPT (Point as) (Point as))
forall (as :: [Nat]) (bs :: [Nat]).
(DF as ~> DF bs)
-> Lens (Point as) (Point as) (Point bs) (Point bs)
toLens a ~> b
DF as ~> DF bs
f)