{-# LANGUAGE CPP #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Pointwise linearization and backpropagation for 'Diff' nets and bodies.
--
-- 'linearizeAt' runs a 'Net (,) (Diff p)' forward at a primal point and builds
-- the transposed net of pointwise pullbacks.  This is the honest reverse-mode
-- gradient net: the graph's structure is burned down into a straight linear
-- (affine) cotangent map.
--
-- For feedback-bearing circuits, use 'linearizeBody' on a 'Body' value.
-- The channel type stays exposed, so star-elimination can operate without
-- recovering evidence from a hidden 'Loop.Knot'.
module Circuit.Diff.Backprop
  ( -- * Pointwise linearization
    linearizeAt,
    fromDiffAt,

    -- * Linearization over the body language
    linearizeBody,
  )
where

import Circuit.Bimonoid
  ( CopyT (..),
    DiscardT (..),
    MergeT (..),
    SigCopy (..),
    SigDiscard (..),
    SigPlus (..),
    SigZero (..),
    ZeroT (..),
  )
import Circuit.Body (Body (..))
import Circuit.Category ((.))
import Circuit.Diff (Diff (..), Diff', runDiff)
import Circuit.Diff.Circuit ()
import Circuit.Net (Net, braid, lift, widen)
import Circuit.Pullback (Pullback (..))
import Circuit.SMC (SMC, SigPar (..), SigSwap (..))
import Circuit.SMC qualified as SMC
import Circuit.Syntax (SigCompose (..), Syntax (..), (:+:) (..))
import Circuit.Tensor (Tensor (..))
import NumHask.Prelude hiding ((.))

-- $setup
-- >>> import Circuit.Category ((.))
-- >>> import Circuit.Diff
-- >>> import Circuit.Bimonoid qualified as Bm
-- >>> import Circuit.Net (Net, lift, widen)
-- >>> import Circuit.Pullback (Pullback (..), evalPullback)
-- >>> import Circuit.SMC qualified as SMC
-- >>> import Circuit.Tensor (Tensor (..))
-- >>> import Prelude hiding (id, (.))

-- | Pointwise linearization: run a 'Net' of 'Diff' primitives forward at @a@
-- and build the transposed net of pullbacks.
--
-- This is the same operation as /backpropagation/, but read forwards
-- through the lens of 'linearize': the graph's structure is burned down
-- into a straight linear (affine) cotangent map.  Looking backwards, the
-- single output cotangent appears to bifurcate and fan out through the
-- wiring.  Propagate fans values out in the forward direction; linearize
-- straightens them into a wire in the reverse direction.
--
-- This is the honest reverse-mode gradient net.  'transpose' alone is
-- only correct for linear nets; for nonlinear 'Diff' primitives the
-- pullback closure depends on the primal point.  'linearizeAt' runs
-- the net, captures each primitive's pullback at the point it saw,
-- and returns both the output value and a 'Net Pullback' whose wires
-- are those pointwise pullbacks composed in reverse order.
--
-- Use 'Circuit.Pullback.evalPullback' to evaluate the resulting net
-- at a single output cotangent.
--
-- >>> let sq = Diff (\x -> (x * x, \d -> 2 * x * d)) :: Diff' Double Double
-- >>> let copyN = lift (Bm.copyT @(,) @Diff' @Double) :: Net (,) Diff' Double (Double, Double)
-- >>> let plusN = lift (Bm.plusT @(,) @Diff' @Double) :: Net (,) Diff' (Double, Double) Double
-- >>> let parN = widen (tensor (SMC.lift sq) (SMC.lift sq)) :: Net (,) Diff' (Double, Double) (Double, Double)
-- >>> let n = plusN . parN . copyN :: Net (,) Diff' Double Double
-- >>> let (y, g) = linearizeAt n 3.0
-- >>> y
-- 18.0
-- >>> evalPullback g 1.0
-- 12.0

-- | Capture the pullback of a 'Diff' primitive at a primal point.
--
-- >>> let d = Diff (\x -> (x * x, \dy -> 2 * x * dy))
-- >>> runPullback (fromDiffAt d 3) 1
-- 6
fromDiffAt :: forall p a b. Diff p a b -> a -> Pullback b a
fromDiffAt :: forall {k} (p :: k) a b. Diff p a b -> a -> Pullback b a
fromDiffAt (Diff a -> (b, b -> a)
f) a
a = (b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback ((b, b -> a) -> b -> a
forall a b. (a, b) -> b
snd (a -> (b, b -> a)
f a
a))
{-# INLINE fromDiffAt #-}

-- | Run a 'Net' of 'Diff' primitives forward and build the transposed
-- pullback net.
--
-- Structural rows ('SigCopy', 'SigPlus', 'SigDiscard', 'SigZero') are
-- converted to point-independent 'lift' pullbacks using the 'Diff'
-- dictionaries the constructors already carry (copy↦plus, plus↦dup,
-- discard↦zero, zero↦discard).
linearizeAt ::
  forall p a b.
  Net (,) (Diff p) a b ->
  a ->
  (b, Net (,) Pullback b a)
linearizeAt :: forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeAt = Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeNet

-- | Pointwise linearization over the free 'Net' language.
linearizeNet ::
  forall p a b.
  Net (,) (Diff p) a b ->
  a ->
  (b, Net (,) Pullback b a)
linearizeNet :: forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeNet Net (,) (Diff p) a b
n a
a = case Net (,) (Diff p) a b
n of
  Lift Diff p a b
d ->
    let (b
y, b -> a
pb) = Diff p a b -> a -> (b, b -> a)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff p a b
d a
a
     in (b
y, Pullback b a -> Net (,) Pullback b a
forall (arr :: * -> * -> *) a b (w :: * -> * -> *).
arr a b -> Net w arr a b
lift ((b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback b -> a
pb))
  Op (:+:)
  SigCompose
  (SigPar (,)
   :+: (SigSwap (,)
        :+: (SigCopy (,)
             :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
op -> case (:+:)
  SigCompose
  (SigPar (,)
   :+: (SigSwap (,)
        :+: (SigCopy (,)
             :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
op of
    L (SigCompose Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  b1
  b
g Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  a
  b1
f) ->
      let (b1
b, Net (,) Pullback b1 a
f') = Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  a
  b1
-> a -> (b1, Net (,) Pullback b1 a)
forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeNet Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  a
  b1
f a
a
          (b
c, Net (,) Pullback b b1
g') = Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  b1
  b
-> b1 -> (b, Net (,) Pullback b b1)
forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeNet Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  b1
  b
g b1
b
       in (b
c, Net (,) Pullback b1 a
f' Net (,) Pullback b1 a
-> Net (,) Pullback b b1 -> Net (,) Pullback b a
forall b c a.
Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  Pullback
  b
  c
-> Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     Pullback
     a
     b
-> Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     Pullback
     a
     c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Net (,) Pullback b b1
g')
    R (L (SigPar Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  a1
  b1
f Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  c
  d
g)) ->
      let (a1
a1, c
a2) = a
a
          (b1
b, Net (,) Pullback b1 a1
f') = Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  a1
  b1
-> a1 -> (b1, Net (,) Pullback b1 a1)
forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeNet Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  a1
  b1
f a1
a1
          (d
d, Net (,) Pullback d c
g') = Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  c
  d
-> c -> (d, Net (,) Pullback d c)
forall {k} (p :: k) a b.
Net (,) (Diff p) a b -> a -> (b, Net (,) Pullback b a)
linearizeNet Syntax
  (SigCompose
   :+: (SigPar (,)
        :+: (SigSwap (,)
             :+: (SigCopy (,)
                  :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
  (Diff p)
  c
  d
g c
a2
       in ((b1
b, d
d), (:+:)
  SigCompose
  (SigPar (,)
   :+: (SigSwap (,)
        :+: (SigCopy (,)
             :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))
  Pullback
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     Pullback)
  b
  a
-> Net (,) Pullback b a
forall (sig :: Sig) (arr :: * -> * -> *) a b.
sig arr (Syntax sig arr) a b -> Syntax sig arr a b
Op ((:+:)
  (SigPar (,))
  (SigSwap (,)
   :+: (SigCopy (,)
        :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))
  Pullback
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     Pullback)
  b
  a
-> (:+:)
     SigCompose
     (SigPar (,)
      :+: (SigSwap (,)
           :+: (SigCopy (,)
                :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,))))))
     Pullback
     (Syntax
        (SigCompose
         :+: (SigPar (,)
              :+: (SigSwap (,)
                   :+: (SigCopy (,)
                        :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
        Pullback)
     b
     a
forall {k} {k1} {k2} {k3} (sig2 :: k -> k1 -> k2 -> k3 -> *)
       (arr :: k) (rec :: k1) (a :: k2) (b :: k3)
       (sig1 :: k -> k1 -> k2 -> k3 -> *).
sig2 arr rec a b -> (:+:) sig1 sig2 arr rec a b
R (SigPar
  (,)
  Pullback
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     Pullback)
  b
  a
-> (:+:)
     (SigPar (,))
     (SigSwap (,)
      :+: (SigCopy (,)
           :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))
     Pullback
     (Syntax
        (SigCompose
         :+: (SigPar (,)
              :+: (SigSwap (,)
                   :+: (SigCopy (,)
                        :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
        Pullback)
     b
     a
forall {k} {k1} {k2} {k3} (sig1 :: k -> k1 -> k2 -> k3 -> *)
       (arr :: k) (rec :: k1) (a :: k2) (b :: k3)
       (sig2 :: k -> k1 -> k2 -> k3 -> *).
sig1 arr rec a b -> (:+:) sig1 sig2 arr rec a b
L (Net (,) Pullback b1 a1
-> Net (,) Pullback d c
-> SigPar
     (,)
     Pullback
     (Syntax
        (SigCompose
         :+: (SigPar (,)
              :+: (SigSwap (,)
                   :+: (SigCopy (,)
                        :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
        Pullback)
     (b1, d)
     (a1, c)
forall {k} (rec :: * -> * -> *) a1 b1 c d (w :: * -> * -> *)
       (arr :: k).
rec a1 b1 -> rec c d -> SigPar w arr rec (w a1 c) (w b1 d)
SigPar Net (,) Pullback b1 a1
f' Net (,) Pullback d c
g'))))
    R (R (L SigSwap
  (,)
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
SigSwap)) ->
      let (a1
u, b1
v) = a
a
       in ((b1
v, a1
u), Net (,) Pullback b a
Net (,) Pullback (b1, a1) (a1, b1)
forall (w :: * -> * -> *) (arr :: * -> * -> *) a b.
Net w arr (w a b) (w b a)
braid)
    R (R (R (L SigCopy
  (,)
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
SigCopy))) ->
      let ((a, a)
out, (a, a) -> a
pb) = Diff p a (a, a) -> a -> ((a, a), (a, a) -> a)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff (forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
CopyT t arr a =>
arr a (t a a)
forall (t :: * -> * -> *) (arr :: * -> * -> *) a.
CopyT t arr a =>
arr a (t a a)
copyT @(,) @(Diff p)) a
a
       in (b
(a, a)
out, Pullback b a -> Net (,) Pullback b a
forall (arr :: * -> * -> *) a b (w :: * -> * -> *).
arr a b -> Net w arr a b
lift ((b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback b -> a
(a, a) -> a
pb))
    R (R (R (R (L SigDiscard
  (,)
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
SigDiscard)))) ->
      let (b
out, b -> a
pb) = Diff p a b -> a -> (b, b -> a)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff (forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
DiscardT t arr a =>
arr a (Unit t)
forall (t :: * -> * -> *) (arr :: * -> * -> *) a.
DiscardT t arr a =>
arr a (Unit t)
discardT @(,) @(Diff p)) a
a
       in (b
out, Pullback b a -> Net (,) Pullback b a
forall (arr :: * -> * -> *) a b (w :: * -> * -> *).
arr a b -> Net w arr a b
lift ((b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback b -> a
pb))
    R (R (R (R (R (L SigPlus
  (,)
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
SigPlus))))) ->
      let (b
out, b -> (b, b)
pb) = Diff p (b, b) b -> (b, b) -> (b, b -> (b, b))
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff (forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
MergeT t arr a =>
arr (t a a) a
forall (t :: * -> * -> *) (arr :: * -> * -> *) a.
MergeT t arr a =>
arr (t a a) a
plusT @(,) @(Diff p)) a
(b, b)
a
       in (b
out, Pullback b a -> Net (,) Pullback b a
forall (arr :: * -> * -> *) a b (w :: * -> * -> *).
arr a b -> Net w arr a b
lift ((b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback b -> a
b -> (b, b)
pb))
    R (R (R (R (R (R SigZero
  (,)
  (Diff p)
  (Syntax
     (SigCompose
      :+: (SigPar (,)
           :+: (SigSwap (,)
                :+: (SigCopy (,)
                     :+: (SigDiscard (,) :+: (SigPlus (,) :+: SigZero (,)))))))
     (Diff p))
  a
  b
SigZero))))) ->
      let (b
out, b -> ()
pb) = Diff p () b -> () -> (b, b -> ())
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff (forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
ZeroT t arr a =>
arr (Unit t) a
forall (t :: * -> * -> *) (arr :: * -> * -> *) a.
ZeroT t arr a =>
arr (Unit t) a
zeroT @(,) @(Diff p)) ()
       in (b
out, Pullback b a -> Net (,) Pullback b a
forall (arr :: * -> * -> *) a b (w :: * -> * -> *).
arr a b -> Net w arr a b
lift ((b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback b -> a
b -> ()
pb))

-- | Pointwise linearization over the core 'Body' language.
--
-- The channel type @s@ stays exposed, so the result can be fed directly to
-- 'Circuit.Diff.Star.solveStarBody' without any coercion.
--
-- /Caveat/: fixpoints are lazy on both passes.  The forward body ties the
-- same lazy knot as 'Trace' @Diff@; the returned body ties the lazy
-- 'Trace' @Pullback@ knot.  For strict carriers with nonzero channel
-- self-coupling, /both/ diverge.  The backward side can be solved in closed
-- form by 'Circuit.Diff.Star.solveStarBody'.
linearizeBody ::
  forall p s a b.
  Body (,) s (Diff p) a b ->
  a ->
  (b, Body (,) s Pullback b a)
linearizeBody :: forall {k} (p :: k) s a b.
Body (,) s (Diff p) a b -> a -> (b, Body (,) s Pullback b a)
linearizeBody (Body Diff p (s, a) (s, b)
f) a
a =
  let ~((s
s, b
b), (s, b) -> (s, a)
pb) = Diff p (s, a) (s, b) -> (s, a) -> ((s, b), (s, b) -> (s, a))
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff p (s, a) (s, b)
f (s
s, a
a)
   in (b
b, Pullback (s, b) (s, a) -> Body (,) s Pullback b a
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (((s, b) -> (s, a)) -> Pullback (s, b) (s, a)
forall b a. (b -> a) -> Pullback b a
Pullback (s, b) -> (s, a)
pb))