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

-- | 'Diff' as a circuits arrow.
--
-- This module gives the instances that turn the differentiable carrier
-- 'Circuit.Diff.Diff' into a 'Circuit.Category.Category' with tracing,
-- channels, strength, tensor products, and bimonoid structure.  Keeping the
-- instances in the same package as the 'Diff' type avoids the orphan
-- instances that would arise if 'circuits-ad' defined them for a carrier
-- living elsewhere.
module Circuit.Diff.Circuit
  ( -- * Re-exports from the carrier
    Diff (..),
    Diff',

    -- * Traced variants
    traceNFrom,
    traceStarFrom,
    traceStar,

    -- * Smoke test
    quadD,
  )
where

import Circuit.Bimonoid (Copy (..), Discard (..), Merge (..), MergeZero, Zero (..))
import Circuit.Bimonoid qualified as CB
import Circuit.Category (Category (..))
import Circuit.Channel (Channel (..), Strength (..), Traced (..))
import Circuit.Diff (Diff (..), Diff', runDiff)
import Circuit.Tensor (Action (..), Tensor (..), Unital (..))
import Data.Bifunctor
import NumHask.Algebra.Additive qualified as NHA
import NumHask.Algebra.Multiplicative qualified as NHM
import NumHask.Algebra.Ring qualified as NHR
import Prelude hiding (id, (.))

-- $setup
-- >>> import Circuit.Bimonoid (Copy (..), Merge (..))
-- >>> import Circuit.Tensor (Action (..), Tensor (..))

-- | 'Circuit.Category.Category' for 'Diff'.
--
-- 'Circuit.Diff still provides 'Control.Category'; circuits needs the local
-- 'Category' with associated 'Ob' (default @()@) so 'Monoidal' / 'Traced' /
-- free 'Trace' folds typecheck after kind-gen.
instance Category (Diff p) where
  id :: forall a. Diff p a a
id = (a -> (a, a -> a)) -> Diff p a a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\a
a -> (a
a, a -> a
forall a. a -> a
forall k (arr :: k -> k -> *) (a :: k). Category arr => arr a a
id))
  Diff b -> (c, c -> b)
f . :: forall b c a. Diff p b c -> Diff p a b -> Diff p a c
. Diff a -> (b, b -> a)
g = (a -> (c, c -> a)) -> Diff p a c
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((a -> (c, c -> a)) -> Diff p a c)
-> (a -> (c, c -> a)) -> Diff p a c
forall a b. (a -> b) -> a -> b
$ \a
a ->
    let (b
b, b -> a
gb) = a -> (b, b -> a)
g a
a
        (c
c, c -> b
fc) = b -> (c, c -> b)
f b
b
     in (c
c, \c
dc -> b -> a
gb (c -> b
fc c
dc))
  {-# INLINE id #-}
  {-# INLINE (.) #-}

-- | 'Trace' for 'Diff with the @(,)@ tensor.
--
-- The forward pass ties the standard lazy knot:
--
-- @
-- let (a, c) = body (a, b) in c
-- @
--
-- The backward pass ties the /same shape/ of knot, transposed.  Given @dc@
-- on the output, the full backward pair @(da, db)@ satisfies a self-referential
-- equation solved by a single lazy binding:
--
-- @
-- let bd = backward (fst bd, dc) in snd bd
-- @
--
-- The knot flows through the /pair/ rather than through the channel cotangent
-- alone — 'backward' is called once, the pair is destructured once, and the
-- shape mirrors the forward knot identically.
--
-- 'pullback' closes over 'backward', which closes over the forward pass's
-- intermediates.  The closure is the tape: no explicit Wengert list is built
-- because GHC's heap holds the graph.  For linear backward maps this is a
-- Neumann series computed lazily; for general maps it is the implicit function
-- theorem as a lazy knot.
instance Traced (,) (Diff p) where
  trace :: forall a b c. Diff p (a, b) (a, c) -> Diff p b c
trace (Diff (a, b) -> ((a, c), (a, c) -> (a, b))
body) = (b -> (c, c -> b)) -> Diff p b c
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((b -> (c, c -> b)) -> Diff p b c)
-> (b -> (c, c -> b)) -> Diff p b c
forall a b. (a -> b) -> a -> b
$ \b
b ->
    let -- Forward: standard lazy knot
        ~((a
a, c
c), (a, c) -> (a, b)
backward) = (a, b) -> ((a, c), (a, c) -> (a, b))
body (a
a, b
b)
        -- Backward: same shape, transposed — knot through the pair
        pullback :: c -> b
pullback c
dc =
          let bd :: (a, b)
bd = (a, c) -> (a, b)
backward ((a, b) -> a
forall a b. (a, b) -> a
fst (a, b)
bd, c
dc)
           in (a, b) -> b
forall a b. (a, b) -> b
snd (a, b)
bd
     in (c
c, c -> b
pullback)

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

-- | Cocartesian channel plumbing for 'Diff.
instance Channel Either (Diff p) where
  assoc :: forall a b c.
Diff p (Either (Either a b) c) (Either a (Either b c))
assoc =
    (Either (Either a b) c
 -> (Either a (Either b c),
     Either a (Either b c) -> Either (Either a b) c))
-> Diff p (Either (Either a b) c) (Either a (Either b c))
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff
      ( \case
          Left (Left a
a) -> (a -> Either a (Either b c)
forall a b. a -> Either a b
Left a
a, \case Left a
da -> Either a b -> Either (Either a b) c
forall a b. a -> Either a b
Left (a -> Either a b
forall a b. a -> Either a b
Left a
da); Right Either b c
_ -> [Char] -> Either (Either a b) c
forall a. HasCallStack => [Char] -> a
error [Char]
"assoc")
          Left (Right b
b) -> (Either b c -> Either a (Either b c)
forall a b. b -> Either a b
Right (b -> Either b c
forall a b. a -> Either a b
Left b
b), \case Right (Left b
db) -> Either a b -> Either (Either a b) c
forall a b. a -> Either a b
Left (b -> Either a b
forall a b. b -> Either a b
Right b
db); Either a (Either b c)
_ -> [Char] -> Either (Either a b) c
forall a. HasCallStack => [Char] -> a
error [Char]
"assoc")
          Right c
c -> (Either b c -> Either a (Either b c)
forall a b. b -> Either a b
Right (c -> Either b c
forall a b. b -> Either a b
Right c
c), \case Right (Right c
dc) -> c -> Either (Either a b) c
forall a b. b -> Either a b
Right c
dc; Either a (Either b c)
_ -> [Char] -> Either (Either a b) c
forall a. HasCallStack => [Char] -> a
error [Char]
"assoc")
      )
  assoc' :: forall a b c.
Diff p (Either a (Either b c)) (Either (Either a b) c)
assoc' =
    (Either a (Either b c)
 -> (Either (Either a b) c,
     Either (Either a b) c -> Either a (Either b c)))
-> Diff p (Either a (Either b c)) (Either (Either a b) c)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff
      ( \case
          Left a
a -> (Either a b -> Either (Either a b) c
forall a b. a -> Either a b
Left (a -> Either a b
forall a b. a -> Either a b
Left a
a), \case Left (Left a
da) -> a -> Either a (Either b c)
forall a b. a -> Either a b
Left a
da; Either (Either a b) c
_ -> [Char] -> Either a (Either b c)
forall a. HasCallStack => [Char] -> a
error [Char]
"assoc'")
          Right (Left b
b) -> (Either a b -> Either (Either a b) c
forall a b. a -> Either a b
Left (b -> Either a b
forall a b. b -> Either a b
Right b
b), \case Left (Right b
db) -> Either b c -> Either a (Either b c)
forall a b. b -> Either a b
Right (b -> Either b c
forall a b. a -> Either a b
Left b
db); Either (Either a b) c
_ -> [Char] -> Either a (Either b c)
forall a. HasCallStack => [Char] -> a
error [Char]
"assoc'")
          Right (Right c
c) -> (c -> Either (Either a b) c
forall a b. b -> Either a b
Right c
c, \case Right c
dc -> Either b c -> Either a (Either b c)
forall a b. b -> Either a b
Right (c -> Either b c
forall a b. b -> Either a b
Right c
dc); Either (Either a b) c
_ -> [Char] -> Either a (Either b c)
forall a. HasCallStack => [Char] -> a
error [Char]
"assoc'")
      )
  slide :: forall a b c.
Diff p (Either a (Either b c)) (Either b (Either a c))
slide =
    (Either a (Either b c)
 -> (Either b (Either a c),
     Either b (Either a c) -> Either a (Either b c)))
-> Diff p (Either a (Either b c)) (Either b (Either a c))
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff
      ( \case
          Left a
a -> (Either a c -> Either b (Either a c)
forall a b. b -> Either a b
Right (a -> Either a c
forall a b. a -> Either a b
Left a
a), \case Right (Left a
da) -> a -> Either a (Either b c)
forall a b. a -> Either a b
Left a
da; Either b (Either a c)
_ -> [Char] -> Either a (Either b c)
forall a. HasCallStack => [Char] -> a
error [Char]
"slide")
          Right (Left b
b) -> (b -> Either b (Either a c)
forall a b. a -> Either a b
Left b
b, \case Left b
db -> Either b c -> Either a (Either b c)
forall a b. b -> Either a b
Right (b -> Either b c
forall a b. a -> Either a b
Left b
db); Either b (Either a c)
_ -> [Char] -> Either a (Either b c)
forall a. HasCallStack => [Char] -> a
error [Char]
"slide")
          Right (Right c
c) -> (Either a c -> Either b (Either a c)
forall a b. b -> Either a b
Right (c -> Either a c
forall a b. b -> Either a b
Right c
c), \case Right (Right c
dc) -> Either b c -> Either a (Either b c)
forall a b. b -> Either a b
Right (c -> Either b c
forall a b. b -> Either a b
Right c
dc); Either b (Either a c)
_ -> [Char] -> Either a (Either b c)
forall a. HasCallStack => [Char] -> a
error [Char]
"slide")
      )

-- | Cartesian tensorial strength for 'Diff.
instance Strength (,) (Diff p) where
  strength :: forall b c a. Diff p b c -> Diff p (a, b) (a, c)
strength (Diff b -> (c, c -> b)
f) = ((a, b) -> ((a, c), (a, c) -> (a, b))) -> Diff p (a, b) (a, c)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (((a, b) -> ((a, c), (a, c) -> (a, b))) -> Diff p (a, b) (a, c))
-> ((a, b) -> ((a, c), (a, c) -> (a, b))) -> Diff p (a, b) (a, c)
forall a b. (a -> b) -> a -> b
$ \(a
a, b
b) ->
    let (c
c, c -> b
back) = b -> (c, c -> b)
f b
b
     in ((a
a, c
c), \(a
da, c
dc) -> (a
da, c -> b
back c
dc))
  {-# INLINE strength #-}

-- | Cocartesian tensorial strength for 'Diff.
instance Strength Either (Diff p) where
  strength :: forall b c a. Diff p b c -> Diff p (Either a b) (Either a c)
strength (Diff b -> (c, c -> b)
f) = (Either a b -> (Either a c, Either a c -> Either a b))
-> Diff p (Either a b) (Either a c)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((Either a b -> (Either a c, Either a c -> Either a b))
 -> Diff p (Either a b) (Either a c))
-> (Either a b -> (Either a c, Either a c -> Either a b))
-> Diff p (Either a b) (Either a c)
forall a b. (a -> b) -> a -> b
$ \case
    Left a
a -> (a -> Either a c
forall a b. a -> Either a b
Left a
a, \case Left a
da -> a -> Either a b
forall a b. a -> Either a b
Left a
da; Right c
_ -> [Char] -> Either a b
forall a. HasCallStack => [Char] -> a
error [Char]
"strength: Left input, Right cotangent")
    Right b
b ->
      let (c
c, c -> b
back) = b -> (c, c -> b)
f b
b
       in (c -> Either a c
forall a b. b -> Either a b
Right c
c, \case Right c
dc -> b -> Either a b
forall a b. b -> Either a b
Right (c -> b
back c
dc); Left a
_ -> [Char] -> Either a b
forall a. HasCallStack => [Char] -> a
error [Char]
"strength: Right input, Left cotangent")
  {-# INLINE strength #-}

-- | Trace for 'Diff with the 'Either' tensor.
--
-- The 'Either' trace is a while-loop: 'Left a' means "iterate again",
-- 'Right c' means "return".  Forward pass runs until the body produces a
-- 'Right', recording each body's pullback.  Backward pass replays those
-- pullbacks in reverse order, propagating the output cotangent back through
-- the iteration chain.
--
-- The number of iterations is treated as locally constant by the derivative:
-- small perturbations of the input do not change the branch sequence.  This
-- is the standard reverse-mode treatment of data-dependent control flow.
--
-- __Proof obligation__ (joins the linearity obligation on the other
-- traces): a cotangent on a sum is represented as the /same/ sum, and
-- its tag must match the primal trajectory — the cotangent space at a
-- point of @Either a c@ is the cotangent space of the branch the point
-- is in.  Every honest pullback maps an output-tagged cotangent to an
-- input-tagged one; the replay errors loudly on any mismatch rather
-- than misreading a dishonest primitive.
instance Traced Either (Diff p) where
  trace :: forall a b c. Diff p (Either a b) (Either a c) -> Diff p b c
trace (Diff Either a b -> (Either a c, Either a c -> Either a b)
body) = (b -> (c, c -> b)) -> Diff p b c
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((b -> (c, c -> b)) -> Diff p b c)
-> (b -> (c, c -> b)) -> Diff p b c
forall a b. (a -> b) -> a -> b
$ \b
b ->
    let -- Forward: iterate, collecting pullbacks in execution order.
        goFwd :: Either a b -> (c, [Either a c -> Either a b])
goFwd Either a b
x =
          let (Either a c
y, Either a c -> Either a b
pb) = Either a b -> (Either a c, Either a c -> Either a b)
body Either a b
x
           in case Either a c
y of
                Right c
c' -> (c
c', [Either a c -> Either a b
pb])
                Left a
a ->
                  let (c
c', [Either a c -> Either a b]
pbs') = Either a b -> (c, [Either a c -> Either a b])
goFwd (a -> Either a b
forall a b. a -> Either a b
Left a
a)
                   in (c
c', Either a c -> Either a b
pb (Either a c -> Either a b)
-> [Either a c -> Either a b] -> [Either a c -> Either a b]
forall a. a -> [a] -> [a]
: [Either a c -> Either a b]
pbs')
        (c
c, [Either a c -> Either a b]
pbs) = Either a b -> (c, [Either a c -> Either a b])
goFwd (b -> Either a b
forall a b. b -> Either a b
Right b
b)
        -- Reverse the tape once; shared across all cotangents.
        rpbs :: [Either a c -> Either a b]
rpbs = [Either a c -> Either a b] -> [Either a c -> Either a b]
forall a. [a] -> [a]
reverse [Either a c -> Either a b]
pbs
        -- Backward: replay pullbacks in reverse order.
        pullback :: c -> b
pullback c
dc =
          case [Either a c -> Either a b]
rpbs of
            [] -> [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Circuit.Trace Diff Either: empty loop (impossible)"
            (Either a c -> Either a b
lastPb : [Either a c -> Either a b]
prevPbs) -> Either a b -> [Either a c -> Either a b] -> b
forall {a} {t} {b}. Either a t -> [Either a b -> Either a t] -> t
goBwd (Either a c -> Either a b
lastPb (c -> Either a c
forall a b. b -> Either a b
Right c
dc)) [Either a c -> Either a b]
prevPbs
          where
            goBwd :: Either a t -> [Either a b -> Either a t] -> t
goBwd (Right t
db) [] = t
db
            goBwd (Left a
_) [] =
              [Char] -> t
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Circuit.Trace Diff Either: final cotangent landed on Left (impossible)"
            goBwd (Left a
da) (Either a b -> Either a t
pb : [Either a b -> Either a t]
pbs') = Either a t -> [Either a b -> Either a t] -> t
goBwd (Either a b -> Either a t
pb (a -> Either a b
forall a b. a -> Either a b
Left a
da)) [Either a b -> Either a t]
pbs'
            goBwd (Right t
_) (Either a b -> Either a t
_ : [Either a b -> Either a t]
_) =
              -- A Right-tagged cotangent with pullbacks still pending means
              -- some pullback at iteration i > 1 claimed its input was the
              -- exit branch — but that iteration's primal input was 'Left'.
              -- Returning here would silently skip the remaining chain rule,
              -- so this is a dishonest primitive, not an early exit.
              [Char] -> t
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Circuit.Trace Diff Either: Right cotangent mid-chain (tag-dishonest pullback)"
     in (c
c, c -> b
pullback)

-- ---------------------------------------------------------------------------
-- Copy / Discard / Merge / Zero for Diff — the bimonoid structure of the
-- differentiable arrow.
--
-- In Diff, the bimonoid is self-dual under differentiation: copy's pullback
-- is plus, discard's pullback is zero, plus's pullback is copy, zero's
-- pullback is discard.  transpose's Copy ↔ Add, Discard ↔ Zero table is
-- not a rule imposed on syntax — it's the instance structure of Diff read
-- off at the semantic level.

-- | Copy in D: the pullback is 'plus' (fan-in on the backward pass).
--
-- >>> import Circuit.Tensor (Action(..))
-- >>> import Circuit.Bimonoid (Copy(..), Merge(..))
-- >>> let (_, pb) = runDiff (copy :: Diff' Int (Int, Int)) 5
-- >>> pb (1, 2)
-- 3
instance (Merge (->) a) => Copy (Diff p) a where
  copy :: Diff p a (a, a)
copy = (a -> ((a, a), (a, a) -> a)) -> Diff p a (a, a)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\a
a -> ((a
a, a
a), (a, a) -> a
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
CB.plus))
  {-# INLINE copy #-}

instance (Zero (->) a) => Discard (Diff p) a where
  discard :: Diff p a ()
discard = (a -> ((), () -> a)) -> Diff p a ()
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (((), () -> a) -> a -> ((), () -> a)
forall a b. a -> b -> a
const ((), \() -> () -> a
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE discard #-}

-- | Add in D: the pullback is 'copy' (fan-out on the backward pass).
--
-- >>> let (_, pb) = runDiff (plus :: Diff' (Int, Int) Int) (3, 4)
-- >>> pb 1
-- (1,1)
instance (Merge (->) a) => Merge (Diff p) a where
  plus :: Diff p (a, a) a
plus = ((a, a) -> (a, a -> (a, a))) -> Diff p (a, a) a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\(a
a, a
b) -> ((a, a) -> a
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
CB.plus (a
a, a
b), \a
d -> (a
d, a
d)))
  {-# INLINE plus #-}

instance (Zero (->) a) => Zero (Diff p) a where
  zero :: Diff p () a
zero = (() -> (a, a -> ())) -> Diff p () a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\() -> (() -> a
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero (), () -> a -> ()
forall a b. a -> b -> a
const ()))
  {-# INLINE zero #-}

-- | Monoidal product for Diff: independent wires, no additive constraint.
--
-- >>> let f = Diff (\x -> (x + 1, \d -> d)) :: Diff' Int Int
-- >>> let g = Diff (\x -> (x * 2, \d -> 2 * d)) :: Diff' Int Int
-- >>> let (y, pb) = runDiff (tensor f g) (3, 4)
-- >>> y
-- (4,8)
-- >>> pb (1, 1)
-- (1,2)
instance Unital (,) (Diff p) where
  unitl :: forall a. Diff p (Unit (,), a) a
unitl = (((), a) -> (a, a -> ((), a))) -> Diff p ((), a) a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\((), a
a) -> (a
a, ((),)))
  {-# INLINE unitl #-}
  unitl' :: forall a. Diff p a (Unit (,), a)
unitl' = (a -> (((), a), ((), a) -> a)) -> Diff p a ((), a)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\a
a -> (((), a
a), \((), a
da) -> a
da))
  {-# INLINE unitl' #-}
  unitr :: forall a. Diff p (a, Unit (,)) a
unitr = ((a, ()) -> (a, a -> (a, ()))) -> Diff p (a, ()) a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\(a
a, ()) -> (a
a, (,())))
  {-# INLINE unitr #-}
  unitr' :: forall a. Diff p a (a, Unit (,))
unitr' = (a -> ((a, ()), (a, ()) -> a)) -> Diff p a (a, ())
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\a
a -> ((a
a, ()), \(a
da, ()) -> a
da))
  {-# INLINE unitr' #-}

instance Tensor (,) (Diff p) where
  tensor :: forall a b c d. Diff p a b -> Diff p c d -> Diff p (a, c) (b, d)
tensor (Diff a -> (b, b -> a)
f) (Diff c -> (d, d -> c)
g) = ((a, c) -> ((b, d), (b, d) -> (a, c))) -> Diff p (a, c) (b, d)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (((a, c) -> ((b, d), (b, d) -> (a, c))) -> Diff p (a, c) (b, d))
-> ((a, c) -> ((b, d), (b, d) -> (a, c))) -> Diff p (a, c) (b, d)
forall a b. (a -> b) -> a -> b
$ \(a
a, c
c) ->
    let (b
b, b -> a
fb) = a -> (b, b -> a)
f a
a; (d
d, d -> c
gd) = c -> (d, d -> c)
g c
c
     in ((b
b, d
d), (b -> a) -> (d -> c) -> (b, d) -> (a, c)
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 b -> a
fb d -> c
gd)
  {-# INLINE tensor #-}

instance Action (,) (Diff p) where
  braid :: forall a b. Diff p (a, b) (b, a)
braid = ((a, b) -> ((b, a), (b, a) -> (a, b))) -> Diff p (a, b) (b, a)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\(a
a, b
b) -> ((b
b, a
a), \(b
db, a
da) -> (a
da, b
db)))
  {-# INLINE braid #-}

-- ---------------------------------------------------------------------------
-- StarSemiring — the principled Neumann index
-- ---------------------------------------------------------------------------

-- | Trace with a closed-form backward pass via the Kleene 'NHR.star'.
--
-- The integer @n@ in 'traceNFrom' truncates a series on /both/ passes.
-- But only the forward fixpoint is genuinely nonlinear; the backward
-- channel equation is affine — calculus promises linearity in
-- cotangents:
--
-- > da = A·da + C·dc        solution:  da = star A · C·dc
--
-- with @star a = one + a·star a@ — the Neumann series as algebra,
-- @1\/(1−a)@ over a field.  Because the pullback is linear, the
-- blocks A and C·dc are /extractable by probing/:
--
-- > backward (dj, dc) = (A·dj + C·dc, B·dj + D·dc)
-- > A    = fst (backward (one,  0))   -- channel self-coupling
-- > C·dc = fst (backward (zero, dc))
--
-- and the trace's pullback is the Schur complement
-- @D·dc + B·star A·C·dc@, recovered with one more probe at the
-- backward fixpoint:
--
-- > db = snd (backward (star A · C·dc, dc))
--
-- (Check: @A·(star A·C·dc) + C·dc = (A·star A + one)·C·dc
-- = star A·C·dc@ — the star law discharges the fixpoint.)
--
-- So: forward still iterates from the caller's seed (no closed form
-- exists for an arbitrary nonlinear fixpoint), but the backward pass
-- is /exact/ in three calls to @backward@ — no Neumann index at all.
-- The @star@ probe is computed once per forward point and shared
-- across all cotangents.
--
-- The closed form is the truncated iteration's limit; pure-Prelude
-- witness at @a = 0.3@, @c = 2@:
--
-- >>> let daIter = iterate (\d -> 0.3 * d + 2.0 * 1.0) 0 !! 200
-- >>> abs (daIter - 1.0 / (1.0 - 0.3) * 2.0) < 1e-12
-- True
--
-- __Caveat__: @numhask@ declares 'NHR.StarSemiring' but ships no
-- instances; the only carriers in the tower are
-- 'NumHask.Free.Carriers.FieldStar' (@star a = recip (1−a)@),
-- @Warshall@, and @MinPlus@.  For bare 'Double' channels and for /vector/
-- channels solved by 'Circuit.Mat.Dense.starMatrix', see
-- @Circuit.Diff.Star@ — the Schur-complement bridge proper.
--
-- __Proof obligation__: the probes assume the pullback is linear.
-- Every honestly-constructed 'Diff primitive satisfies this (a
-- pullback /is/ a linear map); a primitive whose backward closure is
-- affine-with-offset is a bug that this function will silently
-- misread.
traceStarFrom ::
  (NHR.StarSemiring j, MergeZero (->) c) =>
  -- | forward seed
  j ->
  -- | forward iteration count
  Int ->
  Diff p (j, b) (j, c) ->
  Diff p b c
traceStarFrom :: forall {k} j c (p :: k) b.
(StarSemiring j, MergeZero (->) c) =>
j -> Int -> Diff p (j, b) (j, c) -> Diff p b c
traceStarFrom j
x0 Int
n (Diff (j, b) -> ((j, c), (j, c) -> (j, b))
body) = (b -> (c, c -> b)) -> Diff p b c
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((b -> (c, c -> b)) -> Diff p b c)
-> (b -> (c, c -> b)) -> Diff p b c
forall a b. (a -> b) -> a -> b
$ \b
b ->
  let -- Forward: iterate from caller-supplied seed (as 'traceNFrom')
      stepFwd :: j -> j
stepFwd j
x = let ((j
x', c
_), (j, c) -> (j, b)
_) = (j, b) -> ((j, c), (j, c) -> (j, b))
body (j
x, b
b) in j
x'
      a :: j
a = (j -> j) -> j -> [j]
forall a. (a -> a) -> a -> [a]
iterate j -> j
stepFwd j
x0 [j] -> Int -> j
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
      ((j
_, c
c), (j, c) -> (j, b)
backward) = (j, b) -> ((j, c), (j, c) -> (j, b))
body (j
a, b
b)
      -- Probe the channel self-coupling once; star it in closed form
      aStar :: j
aStar = j -> j
forall a. StarSemiring a => a -> a
NHR.star ((j, b) -> j
forall a b. (a, b) -> a
fst ((j, c) -> (j, b)
backward (j
forall a. Multiplicative a => a
NHM.one, () -> c
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ())))
      -- Backward: exact in two more probes — no iteration
      pullback :: c -> b
pullback c
dc =
        let cdc :: j
cdc = (j, b) -> j
forall a b. (a, b) -> a
fst ((j, c) -> (j, b)
backward (j
forall a. Additive a => a
NHA.zero, c
dc))
         in (j, b) -> b
forall a b. (a, b) -> b
snd ((j, c) -> (j, b)
backward (j
aStar j -> j -> j
forall a. Multiplicative a => a -> a -> a
NHM.* j
cdc, c
dc))
   in (c
c, c -> b
pullback)

-- | Trace via the Kleene star — the execution formula, lazy form.
--
-- For a knot body with channel self-coupling block A and cross-blocks
-- B, C, D, the trace is the Schur complement:
--
-- > traceStar f = D + B · star A · C
--
-- The lazy 'trace' instance for 'Diff computes exactly this via a
-- lazy fixpoint rather than closed form, so this alias is definable
-- without using 'NHR.star' at all.  Note that @numhask@ ships no
-- 'NHR.StarSemiring' instances, so for concrete carriers prefer
-- 'traceStarFrom' (scalar channel, closed-form backward) or
-- @Circuit.Diff.Star.traceStarMatrix@ (vector channel, solved by
-- 'Circuit.Mat.Dense.starMatrix' — the bridge made literal).
traceStar :: Diff p (j, b) (j, c) -> Diff p b c
traceStar :: forall k (p :: k) a b c. Diff p (a, b) (a, c) -> Diff p b c
traceStar = Diff p (j, b) (j, c) -> Diff p b c
forall a b c. Diff p (a, b) (a, c) -> Diff p b c
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k)
       (b :: k) (c :: k).
Traced t arr =>
arr (t a b) (t a c) -> arr b c
trace

-- | Iterated trace for strict carriers.
--
-- The lazy 'trace' diverges on strict cotangent types ('Double', etc.) when
-- the feedback channel has nonzero self-coupling (@∂a_out\/∂a_in ≠ 0@).
-- 'traceNFrom' replaces the lazy knot with truncated fixed-point iteration.
--
--   * __Forward__ — iterate from caller-supplied seed @x0@, N steps.
--     There is no canonical seed for the forward pass (the fixpoint is
--     arbitrary nonlinear), so the caller provides one.
--
--   * __Backward__ — iterate from 'zero', N steps, extract 'snd' once.
--     The backward equation is guaranteed affine (calculus promises
--     linearity in cotangents), so 'zero' is the principled seed.
--     The Neumann summation happens /inside/ the iteration — no
--     double-counting, and 'plus' retreats to where the theory says
--     it lives: inside prims and 'Copy'.
--
-- Lives beside the lawful-but-lazy instance, not replacing it.
traceNFrom ::
  (MergeZero (->) a) =>
  a ->
  Int ->
  Diff p (a, b) (a, c) ->
  Diff p b c
traceNFrom :: forall {k} a (p :: k) b c.
MergeZero (->) a =>
a -> Int -> Diff p (a, b) (a, c) -> Diff p b c
traceNFrom a
x0 Int
n (Diff (a, b) -> ((a, c), (a, c) -> (a, b))
body) = (b -> (c, c -> b)) -> Diff p b c
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((b -> (c, c -> b)) -> Diff p b c)
-> (b -> (c, c -> b)) -> Diff p b c
forall a b. (a -> b) -> a -> b
$ \b
b ->
  let -- Forward: iterate from caller-supplied seed
      stepFwd :: a -> a
stepFwd a
x = let ((a
x', c
_), (a, c) -> (a, b)
_) = (a, b) -> ((a, c), (a, c) -> (a, b))
body (a
x, b
b) in a
x'
      a :: a
a = (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate a -> a
stepFwd a
x0 [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
      ((a
_, c
c), (a, c) -> (a, b)
backward) = (a, b) -> ((a, c), (a, c) -> (a, b))
body (a
a, b
b)
      -- Backward: iterate from zero, extract result once
      pullback :: c -> b
pullback c
dc =
        let stepBwd :: a -> a
stepBwd a
d = (a, b) -> a
forall a b. (a, b) -> a
fst ((a, c) -> (a, b)
backward (a
d, c
dc))
            da :: a
da = (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate a -> a
stepBwd (() -> a
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()) [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
         in (a, b) -> b
forall a b. (a, b) -> b
snd ((a, c) -> (a, b)
backward (a
da, c
dc))
   in (c
c, c -> b
pullback)

-- ---------------------------------------------------------------------------
-- Smoke test: quadratic — the term that was impossible in Stage 1
-- ---------------------------------------------------------------------------

-- | @2x² + 3x + 5@ built from 'tensor', 'dup', and 'plus' on the @Diff@ arrow.
-- No 'Net' needed — the instances are the denotations the rows will realise to.
--
-- The gradient is @4x + 3@, so at @x = 1@: value 10, gradient 7.
--
-- >>> let (y, pb) = runDiff quadD 1.0
-- >>> y
-- 10.0
-- >>> pb 1.0
-- 7.0
quadD :: Diff p Double Double
quadD :: forall {k} (p :: k). Diff p Double Double
quadD = Diff p (Double, Double) Double
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
CB.plus Diff p (Double, Double) Double
-> Diff p (Double, Double) (Double, Double)
-> Diff p (Double, Double) Double
forall b c a. Diff p b c -> Diff p a b -> Diff p a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Diff p Double Double
-> Diff p Double Double -> Diff p (Double, Double) (Double, Double)
forall a b c d. Diff p a b -> Diff p c d -> Diff p (a, c) (b, d)
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k)
       (b :: k) (c :: k) (d :: k).
Tensor t arr =>
arr a b -> arr c d -> arr (t a c) (t b d)
tensor Diff p Double Double
forall {k} (p :: k). Diff p Double Double
sq Diff p Double Double
forall {k} (p :: k). Diff p Double Double
lin Diff p (Double, Double) Double
-> Diff p Double (Double, Double) -> Diff p Double Double
forall b c a. Diff p b c -> Diff p a b -> Diff p a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Diff p Double (Double, Double)
forall (arr :: * -> * -> *) a. Copy arr a => arr a (a, a)
CB.copy
  where
    sq :: Diff p Double Double
sq = (Double -> (Double, Double -> Double)) -> Diff p Double Double
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\Double
x -> (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x, \Double
d -> Double
4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d))
    lin :: Diff p Double Double
lin = (Double -> (Double, Double -> Double)) -> Diff p Double Double
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\Double
x -> (Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
5, (Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
*)))