{-# LANGUAGE CPP #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Linear cotangent maps — the base arrow for reverse-mode gradients.
--
-- @Pullback b a@ is a linear map @b -> a@ read as an arrow from @b@
-- (output cotangent) to @a@ (input cotangent).  Composition is plain
-- function composition — the /reversal/ is not in this category, it is
-- in how a net is transposed so that output cotangents flow back to input
-- cotangents.  Within an arrow the chain rule is then just @(.)@.
--
-- This is the linear semantics behind reverse-mode automatic
-- differentiation: a 'Net' whose wires carry pullbacks rather than smooth
-- maps, avoiding the second-derivative confusion that comes from trying to
-- compose 'Diff' arrows directly.
module Circuit.Pullback
  ( -- * Linear cotangent arrow
    Pullback (..),

    -- * Running a pullback net
    evalPullback,
  )
where

import Circuit.Bimonoid (Copy (..), Discard (..), Merge (..), Zero (..))
import Circuit.Category (Category (..))
import Circuit.Channel (Channel (..), Strength (..), Traced (..))
import Circuit.Layer (run)
import Circuit.Net (Net)
import Circuit.Tensor (Action (..), Tensor (..), Unital (..))
import Data.Bifunctor
import Prelude hiding (id, (.))

-- $setup
-- >>> import Circuit.Category (Category (..))
-- >>> import Circuit.Bimonoid (Copy (..), Discard (..), Merge (..), Zero (..))
-- >>> import Circuit.Channel (Channel (..), Strength (..), Traced (..))
-- >>> import Circuit.Tensor (Action (..), Tensor (..), Unital (..))
-- >>> import Circuit.Net (Net)
-- >>> import Prelude hiding (id, (.))

-- | A linear map from output cotangents to input cotangents, read as
-- an arrow @b -> a@.
--
-- >>> let pb = Pullback (*2) :: Pullback Double Double
-- >>> runPullback pb 3
-- 6.0
newtype Pullback b a = Pullback
  { -- | Apply the pullback to an output cotangent.
    forall b a. Pullback b a -> b -> a
runPullback :: b -> a
  }

instance Category Pullback where
  id :: forall a. Pullback a a
id = (a -> a) -> Pullback a a
forall b a. (b -> a) -> Pullback b a
Pullback a -> a
forall a. a -> a
forall k (arr :: k -> k -> *) (a :: k). Category arr => arr a a
id
  Pullback b -> c
g . :: forall b c a. Pullback b c -> Pullback a b -> Pullback a c
. Pullback a -> b
f = (a -> c) -> Pullback a c
forall b a. (b -> a) -> Pullback b a
Pullback (\a
x -> b -> c
g (a -> b
f a
x))
  {-# INLINE id #-}
  {-# INLINE (.) #-}

-- | Parallel composition pairs pullbacks independently; 'braid' swaps
-- the two cotangents.
--
-- >>> let f = Pullback (+1) :: Pullback Int Int
-- >>> let g = Pullback (*2) :: Pullback Int Int
-- >>> runPullback (tensor f g) (3, 4)
-- (4,8)
instance Unital (,) Pullback where
  unitl :: forall a. Pullback (Unit (,), a) a
unitl = (((), a) -> a) -> Pullback ((), a) a
forall b a. (b -> a) -> Pullback b a
Pullback ((), a) -> a
forall a b. (a, b) -> b
snd
  {-# INLINE unitl #-}
  unitl' :: forall a. Pullback a (Unit (,), a)
unitl' = (a -> ((), a)) -> Pullback a ((), a)
forall b a. (b -> a) -> Pullback b a
Pullback ((),)
  {-# INLINE unitl' #-}
  unitr :: forall a. Pullback (a, Unit (,)) a
unitr = ((a, ()) -> a) -> Pullback (a, ()) a
forall b a. (b -> a) -> Pullback b a
Pullback (a, ()) -> a
forall a b. (a, b) -> a
fst
  {-# INLINE unitr #-}
  unitr' :: forall a. Pullback a (a, Unit (,))
unitr' = (a -> (a, ())) -> Pullback a (a, ())
forall b a. (b -> a) -> Pullback b a
Pullback (,())
  {-# INLINE unitr' #-}

instance Tensor (,) Pullback where
  tensor :: forall a b c d.
Pullback a b -> Pullback c d -> Pullback (a, c) (b, d)
tensor (Pullback a -> b
f) (Pullback c -> d
g) = ((a, c) -> (b, d)) -> Pullback (a, c) (b, d)
forall b a. (b -> a) -> Pullback b a
Pullback ((a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
Data.Bifunctor.bimap a -> b
f c -> d
g)
  {-# INLINE tensor #-}

instance Action (,) Pullback where
  braid :: forall a b. Pullback (a, b) (b, a)
braid = ((a, b) -> (b, a)) -> Pullback (a, b) (b, a)
forall b a. (b -> a) -> Pullback b a
Pullback (\(a
b, b
a) -> (b
a, a
b))
  {-# INLINE braid #-}

instance Strength (,) Pullback where
  strength :: forall b c a. Pullback b c -> Pullback (a, b) (a, c)
strength (Pullback b -> c
f) = ((a, b) -> (a, c)) -> Pullback (a, b) (a, c)
forall b a. (b -> a) -> Pullback b a
Pullback (\(a
a, b
b) -> (a
a, b -> c
f b
b))
  {-# INLINE strength #-}

-- | The cartesian trace for pullbacks.
--
-- The body is a linear map @f :: (x, c) -> (x, b)@.  The traced
-- pullback @c -> b@ solves the affine feedback equation in cotangent
-- space:
--
-- > (dx, db) = f (dx, dc)
--
-- solved by the same lazy knot that a differentiable arrow uses.  For
-- strict carriers with nonzero channel self-coupling this diverges,
-- exactly as the lazy differentiable trace does.  Unlike the differentiable
-- case, though, the equation here is /always affine/ — 'Pullback' arrows are
-- linear by construction — so a knot over a star-semiring carrier can be
-- eliminated outright rather than iterated.
--
-- >>> let body = Pullback (\(dx', dc) -> (2.0 * dc, dx')) :: Pullback (Double, Double) (Double, Double)
-- >>> runPullback (trace body) 1.0
-- 2.0
instance Traced (,) Pullback where
  trace :: forall a b c. Pullback (a, b) (a, c) -> Pullback b c
trace (Pullback (a, b) -> (a, c)
f) = (b -> c) -> Pullback b c
forall b a. (b -> a) -> Pullback b a
Pullback ((b -> c) -> Pullback b c) -> (b -> c) -> Pullback b c
forall a b. (a -> b) -> a -> b
$ \b
dc ->
    let ~(a
dx, c
db) = (a, b) -> (a, c)
f (a
dx, b
dc)
     in c
db
  {-# INLINE trace #-}

-- | Pullback-instance of the comonoid structure.
--
-- Copy's pullback is addition; discard's pullback is the zero
-- cotangent.  These are not used by the transposition step of reverse-mode
-- AD (which encodes structural rows as @Lift@s to avoid channel-type
-- constraints), but they make 'Pullback' a full bimonoid carrier.
--
-- >>> runPullback (copy :: Pullback Int (Int, Int)) 3
-- (3,3)
-- >>> runPullback (discard :: Pullback Int ()) 5
-- ()
--
-- NOTE: neither method here uses an @Additive (->) a@ constraint — copying
-- and discarding are linear as they stand.  If the class head permits,
-- drop the constraint; keeping a stray @Additive@ reads as "addition
-- happens in this instance", which is the confusion the paragraph above
-- tries to dispel.
instance Copy Pullback a where
  copy :: Pullback a (a, a)
copy = (a -> (a, a)) -> Pullback a (a, a)
forall b a. (b -> a) -> Pullback b a
Pullback (\a
b -> (a
b, a
b))
  {-# INLINE copy #-}

instance Discard Pullback a where
  discard :: Pullback a ()
discard = (a -> ()) -> Pullback a ()
forall b a. (b -> a) -> Pullback b a
Pullback (() -> a -> ()
forall a b. a -> b -> a
const ())
  {-# INLINE discard #-}

-- | Pullback-instance of the additive/monoid structure.
--
-- Addition's pullback is copying; zero's pullback is discarding.
--
-- >>> runPullback (plus :: Pullback (Int, Int) Int) (1, 2)
-- 3
-- >>> runPullback (zero :: Pullback () Int) ()
-- 0
instance (Merge (->) a) => Merge Pullback a where
  plus :: Pullback (a, a) a
plus = ((a, a) -> a) -> Pullback (a, a) a
forall b a. (b -> a) -> Pullback b a
Pullback (\(a
b1, a
b2) -> (a, a) -> a
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
plus (a
b1, a
b2))
  {-# INLINE plus #-}

instance (Zero (->) a) => Zero Pullback a where
  zero :: Pullback () a
zero = (() -> a) -> Pullback () a
forall b a. (b -> a) -> Pullback b a
Pullback (\() -> () -> a
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
zero ())
  {-# INLINE zero #-}

-- | Evaluate a pullback net at a single output cotangent.
--
-- This is the one-shot reverse pass: the net was built by transposing a
-- smooth net, and applying it to a cotangent @db@ yields the input
-- cotangent @da@.
evalPullback :: Net (,) Pullback b a -> b -> a
evalPullback :: forall b a. Net (,) Pullback b a -> b -> a
evalPullback Net (,) Pullback b a
n = Pullback b a -> b -> a
forall b a. Pullback b a -> b -> a
runPullback (Net (,) Pullback b a -> Pullback b a
forall (arr :: * -> * -> *) a b.
(Run
   (Syntax
      (SigCompose
       :+: (SigPar (,)
            :+: (SigSwap (,)
                 :+: (SigCopy (,)
                      :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))))
   arr,
 Law
   (Syntax
      (SigCompose
       :+: (SigPar (,)
            :+: (SigSwap (,)
                 :+: (SigCopy (,)
                      :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))))
   arr,
 Bind
   (Syntax
      (SigCompose
       :+: (SigPar (,)
            :+: (SigSwap (,)
                 :+: (SigCopy (,)
                      :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))))
   arr) =>
Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  arr
  a
  b
-> arr a b
forall (f :: (* -> * -> *) -> * -> * -> *) (arr :: * -> * -> *) a
       b.
(Layer f, Run f arr, Law f arr, Bind f arr) =>
f arr a b -> arr a b
run Net (,) Pullback b a
n)
{-# INLINE evalPullback #-}

-- | Cartesian channel plumbing for pullbacks.
instance Channel (,) Pullback where
  assoc :: forall a b c. Pullback ((a, b), c) (a, (b, c))
assoc = (((a, b), c) -> (a, (b, c))) -> Pullback ((a, b), c) (a, (b, c))
forall b a. (b -> a) -> Pullback b a
Pullback (\((a
s, b
s'), c
x) -> (a
s, (b
s', c
x)))
  assoc' :: forall a b c. Pullback (a, (b, c)) ((a, b), c)
assoc' = ((a, (b, c)) -> ((a, b), c)) -> Pullback (a, (b, c)) ((a, b), c)
forall b a. (b -> a) -> Pullback b a
Pullback (\(a
s, (b
s', c
x)) -> ((a
s, b
s'), c
x))
  slide :: forall a b c. Pullback (a, (b, c)) (b, (a, c))
slide = ((a, (b, c)) -> (b, (a, c))) -> Pullback (a, (b, c)) (b, (a, c))
forall b a. (b -> a) -> Pullback b a
Pullback (\(a
s, (b
s', c
x)) -> (b
s', (a
s, c
x)))