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

-- | Parameterised reverse-mode automatic differentiation.
--
-- 'DiffP' is a differentiable operation with parameters @p@, input @a@ and
-- output @b@.  Running forward produces the output; the backward pass, given
-- a cotangent on the output, produces the input cotangent and parameter
-- gradients.
--
-- This is the same shape as 'Circuit.Diff.Diff' from @circuits-ad@, but with
-- an explicit parameter carrier @p@.  When @p = ()@ and the parameter
-- gradient is discarded, 'DiffP' reduces to 'Diff' (see 'toParam' and
-- 'fromParam').
--
-- The representation is the denotation @p -> a -> (b, b -> (a, p))@ rather
-- than a pair of separate forward/backward fields, so it stays aligned with
-- the monomial/lens reading of AD.
module Circuit.Diff.Param
  ( -- * Parameterised differentiable arrow
    DiffP (..),

    -- * Primitive contract
    TensorPrim (..),
    fromPrim,
    toPrim,

    -- * Wiring helpers
    residual,
    splitP,
    joinP,

    -- * Relationship to the phantom-tagged Diff arrow
    toParam,
    fromParam,
  )
where

import Circuit.Bimonoid (Copy (..), Discard (..), Merge (..), MergeZero, Zero (..))
import Circuit.Bimonoid qualified as CB
import Circuit.Category (Category (..))
import Circuit.Channel (Channel (..))
import Circuit.Diff (Diff (..), runDiff)
import Circuit.Tensor (Action (..), Tensor (..), Unital (..))
import Prelude hiding (id, (.))

-- $setup
-- >>> import Circuit.Diff.Param
-- >>> import Circuit.Category (Category (..))
-- >>> import Circuit.Bimonoid (Copy (..), Discard (..), Merge (..), MergeZero, Zero (..))
-- >>> import Circuit.Tensor (Action (..), Tensor (..))
-- >>> import Circuit.Diff (Diff (..), Diff', runDiff)
-- >>> import Prelude hiding (id, (.))

----------------------------------------------------------------------
-- Core type
----------------------------------------------------------------------

-- | A differentiable operation with parameters @p@, input @a@ and output @b@.
--
-- The forward pass closes over the parameter @p@ and input @a@, producing an
-- output @b@ and a pullback.  The pullback maps an output cotangent @db@ to
-- an input cotangent @da@ and a parameter gradient @dp@.
--
-- This is a point-dependent lens: the parameter is part of the position,
-- not an extra layer.
newtype DiffP p a b = DiffP
  { forall p a b. DiffP p a b -> p -> a -> (b, b -> (a, p))
runDiffP :: p -> a -> (b, b -> (a, p))
  }

-- | A primitive contract with separate forward and backward fields, easier
-- to read and write for dense linear-algebra primitives.
data TensorPrim p a b = TensorPrim
  { forall p a b. TensorPrim p a b -> p -> a -> b
primForward :: p -> a -> b,
    forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward :: p -> a -> b -> (a, p)
  }

-- | Convert a 'TensorPrim' into a 'DiffP'.
fromPrim :: TensorPrim p a b -> DiffP p a b
fromPrim :: forall p a b. TensorPrim p a b -> DiffP p a b
fromPrim (TensorPrim p -> a -> b
f p -> a -> b -> (a, p)
b) = (p -> a -> (b, b -> (a, p))) -> DiffP p a b
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> (b, b -> (a, p))) -> DiffP p a b)
-> (p -> a -> (b, b -> (a, p))) -> DiffP p a b
forall a b. (a -> b) -> a -> b
$ \p
p a
a -> (p -> a -> b
f p
p a
a, p -> a -> b -> (a, p)
b p
p a
a)
{-# INLINE fromPrim #-}

-- | Convert a 'DiffP' into a 'TensorPrim'.
toPrim :: DiffP p a b -> TensorPrim p a b
toPrim :: forall p a b. DiffP p a b -> TensorPrim p a b
toPrim (DiffP p -> a -> (b, b -> (a, p))
f) = (p -> a -> b) -> (p -> a -> b -> (a, p)) -> TensorPrim p a b
forall p a b.
(p -> a -> b) -> (p -> a -> b -> (a, p)) -> TensorPrim p a b
TensorPrim (\p
p a
a -> (b, b -> (a, p)) -> b
forall a b. (a, b) -> a
fst (p -> a -> (b, b -> (a, p))
f p
p a
a)) (\p
p a
a b
db -> (b, b -> (a, p)) -> b -> (a, p)
forall a b. (a, b) -> b
snd (p -> a -> (b, b -> (a, p))
f p
p a
a) b
db)
{-# INLINE toPrim #-}

----------------------------------------------------------------------
-- Category
----------------------------------------------------------------------

-- | Sequential composition threads the same parameter @p@ through both
-- arrows and combines the parameter gradients with the monoid structure on
-- @p@.
--
-- Identity returns a zero parameter gradient.  This instance needs
-- @MergeZero (->) p@ because the parameter type must supply both @zero@ (for
-- 'id') and @plus@ (for composing gradients).
--
-- >>> let inc = DiffP (\() x -> (x + 1, \dy -> (dy, ()))) :: DiffP () Int Int
-- >>> let dbl = DiffP (\() x -> (2 * x, \dy -> (2 * dy, ()))) :: DiffP () Int Int
-- >>> let (y, pb) = runDiffP (dbl . inc) () 3
-- >>> y
-- 8
-- >>> pb 1
-- (2,())
--
-- Parameter gradients from both composed arrows accumulate (not just the
-- outer one).  @addParam@ contributes @dpG = dy@; @dblParam@ contributes
-- @dpF = dc@; composition must sum them via 'CB.plus'.
--
-- >>> let addParam = DiffP (\p x -> (x + p, \dy -> (dy, dy))) :: DiffP Int Int Int
-- >>> let dblParam = DiffP (\p b -> (2 * b, \dc -> (2 * dc, dc))) :: DiffP Int Int Int
-- >>> let (y, pb) = runDiffP (dblParam . addParam) 10 3
-- >>> y
-- 26
-- >>> pb 1
-- (2,3)
instance (MergeZero (->) p) => Category (DiffP p) where
  id :: forall a. DiffP p a a
id = (p -> a -> (a, a -> (a, p))) -> DiffP p a a
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> (a, a -> (a, p))) -> DiffP p a a)
-> (p -> a -> (a, a -> (a, p))) -> DiffP p a a
forall a b. (a -> b) -> a -> b
$ \p
_ a
a -> (a
a, (,() -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE id #-}

  DiffP p -> b -> (c, c -> (b, p))
f . :: forall b c a. DiffP p b c -> DiffP p a b -> DiffP p a c
. DiffP p -> a -> (b, b -> (a, p))
g = (p -> a -> (c, c -> (a, p))) -> DiffP p a c
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> (c, c -> (a, p))) -> DiffP p a c)
-> (p -> a -> (c, c -> (a, p))) -> DiffP p a c
forall a b. (a -> b) -> a -> b
$ \p
p a
a ->
    let (b
b, b -> (a, p)
gBack) = p -> a -> (b, b -> (a, p))
g p
p a
a
        (c
c, c -> (b, p)
fBack) = p -> b -> (c, c -> (b, p))
f p
p b
b
     in ( c
c,
          \c
dc ->
            let (b
db, p
dpF) = c -> (b, p)
fBack c
dc
                (a
da, p
dpG) = b -> (a, p)
gBack b
db
             in (a
da, (p, p) -> p
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
CB.plus (p
dpF, p
dpG))
        )
  {-# INLINE (.) #-}

----------------------------------------------------------------------
-- Monoidal structure on objects
----------------------------------------------------------------------

-- | Structural associator for nested pairs.  No parameters are touched.
--
-- The superclass 'Category (DiffP p)' requires 'MergeZero (->) p', so the
-- instance carries the same constraint even though the structural maps
-- themselves ignore the parameter.
instance (MergeZero (->) p) => Channel (,) (DiffP p) where
  assoc :: forall a b c. DiffP p ((a, b), c) (a, (b, c))
assoc = (p
 -> ((a, b), c) -> ((a, (b, c)), (a, (b, c)) -> (((a, b), c), p)))
-> DiffP p ((a, b), c) (a, (b, c))
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p
  -> ((a, b), c) -> ((a, (b, c)), (a, (b, c)) -> (((a, b), c), p)))
 -> DiffP p ((a, b), c) (a, (b, c)))
-> (p
    -> ((a, b), c) -> ((a, (b, c)), (a, (b, c)) -> (((a, b), c), p)))
-> DiffP p ((a, b), c) (a, (b, c))
forall a b. (a -> b) -> a -> b
$ \p
_ ((a
a, b
b), c
c) -> ((a
a, (b
b, c
c)), \(a
da, (b
db, c
dc)) -> (((a
da, b
db), c
dc), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE assoc #-}

  assoc' :: forall a b c. DiffP p (a, (b, c)) ((a, b), c)
assoc' = (p
 -> (a, (b, c)) -> (((a, b), c), ((a, b), c) -> ((a, (b, c)), p)))
-> DiffP p (a, (b, c)) ((a, b), c)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p
  -> (a, (b, c)) -> (((a, b), c), ((a, b), c) -> ((a, (b, c)), p)))
 -> DiffP p (a, (b, c)) ((a, b), c))
-> (p
    -> (a, (b, c)) -> (((a, b), c), ((a, b), c) -> ((a, (b, c)), p)))
-> DiffP p (a, (b, c)) ((a, b), c)
forall a b. (a -> b) -> a -> b
$ \p
_ (a
a, (b
b, c
c)) -> (((a
a, b
b), c
c), \((a
da, b
db), c
dc) -> ((a
da, (b
db, c
dc)), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE assoc' #-}

  slide :: forall a b c. DiffP p (a, (b, c)) (b, (a, c))
slide = (p
 -> (a, (b, c)) -> ((b, (a, c)), (b, (a, c)) -> ((a, (b, c)), p)))
-> DiffP p (a, (b, c)) (b, (a, c))
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p
  -> (a, (b, c)) -> ((b, (a, c)), (b, (a, c)) -> ((a, (b, c)), p)))
 -> DiffP p (a, (b, c)) (b, (a, c)))
-> (p
    -> (a, (b, c)) -> ((b, (a, c)), (b, (a, c)) -> ((a, (b, c)), p)))
-> DiffP p (a, (b, c)) (b, (a, c))
forall a b. (a -> b) -> a -> b
$ \p
_ (a
a, (b
b, c
c)) -> ((b
b, (a
a, c
c)), \(b
db, (a
da, c
dc)) -> ((a
da, (b
db, c
dc)), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE slide #-}

----------------------------------------------------------------------
-- Monoidal product of morphisms
----------------------------------------------------------------------

-- | Parallel composition runs both arrows with the same parameter @p@ and
-- combines their parameter gradients.
--
-- The parameter is treated as an input value, not a state, so it can be
-- passed to both branches without a comonoid structure.  Only the output
-- gradients need @plus@.
--
-- >>> let inc = DiffP (\() x -> (x + 1, \dy -> (dy, ()))) :: DiffP () Int Int
-- >>> let dbl = DiffP (\() x -> (2 * x, \dy -> (2 * dy, ()))) :: DiffP () Int Int
-- >>> let (y, pb) = runDiffP (tensor inc dbl) () (3, 4)
-- >>> y
-- (4,8)
-- >>> pb (1, 1)
-- ((1,2),())
instance (MergeZero (->) p) => Unital (,) (DiffP p) where
  unitl :: forall a. DiffP p (Unit (,), a) a
unitl = (p -> (Unit (,), a) -> (a, a -> ((Unit (,), a), p)))
-> DiffP p (Unit (,), a) a
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> (Unit (,), a) -> (a, a -> ((Unit (,), a), p)))
 -> DiffP p (Unit (,), a) a)
-> (p -> (Unit (,), a) -> (a, a -> ((Unit (,), a), p)))
-> DiffP p (Unit (,), a) a
forall a b. (a -> b) -> a -> b
$ \p
_ ((), a
a) -> (a
a, \a
da -> (((), a
da), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE unitl #-}
  unitl' :: forall a. DiffP p a (Unit (,), a)
unitl' = (p -> a -> ((Unit (,), a), (Unit (,), a) -> (a, p)))
-> DiffP p a (Unit (,), a)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> ((Unit (,), a), (Unit (,), a) -> (a, p)))
 -> DiffP p a (Unit (,), a))
-> (p -> a -> ((Unit (,), a), (Unit (,), a) -> (a, p)))
-> DiffP p a (Unit (,), a)
forall a b. (a -> b) -> a -> b
$ \p
_ a
a -> (((), a
a), \((), a
da) -> (a
da, () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE unitl' #-}
  unitr :: forall a. DiffP p (a, Unit (,)) a
unitr = (p -> (a, Unit (,)) -> (a, a -> ((a, Unit (,)), p)))
-> DiffP p (a, Unit (,)) a
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> (a, Unit (,)) -> (a, a -> ((a, Unit (,)), p)))
 -> DiffP p (a, Unit (,)) a)
-> (p -> (a, Unit (,)) -> (a, a -> ((a, Unit (,)), p)))
-> DiffP p (a, Unit (,)) a
forall a b. (a -> b) -> a -> b
$ \p
_ (a
a, ()) -> (a
a, \a
da -> ((a
da, ()), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE unitr #-}
  unitr' :: forall a. DiffP p a (a, Unit (,))
unitr' = (p -> a -> ((a, Unit (,)), (a, Unit (,)) -> (a, p)))
-> DiffP p a (a, Unit (,))
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> ((a, Unit (,)), (a, Unit (,)) -> (a, p)))
 -> DiffP p a (a, Unit (,)))
-> (p -> a -> ((a, Unit (,)), (a, Unit (,)) -> (a, p)))
-> DiffP p a (a, Unit (,))
forall a b. (a -> b) -> a -> b
$ \p
_ a
a -> ((a
a, ()), \(a
da, ()) -> (a
da, () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE unitr' #-}

instance (MergeZero (->) p) => Tensor (,) (DiffP p) where
  tensor :: forall a b c d. DiffP p a b -> DiffP p c d -> DiffP p (a, c) (b, d)
tensor (DiffP p -> a -> (b, b -> (a, p))
f) (DiffP p -> c -> (d, d -> (c, p))
g) = (p -> (a, c) -> ((b, d), (b, d) -> ((a, c), p)))
-> DiffP p (a, c) (b, d)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> (a, c) -> ((b, d), (b, d) -> ((a, c), p)))
 -> DiffP p (a, c) (b, d))
-> (p -> (a, c) -> ((b, d), (b, d) -> ((a, c), p)))
-> DiffP p (a, c) (b, d)
forall a b. (a -> b) -> a -> b
$ \p
p (a
a, c
c) ->
    let (b
b, b -> (a, p)
fBack) = p -> a -> (b, b -> (a, p))
f p
p a
a
        (d
d, d -> (c, p)
gBack) = p -> c -> (d, d -> (c, p))
g p
p c
c
     in ( (b
b, d
d),
          \(b
db, d
dd) ->
            let (a
da, p
dpF) = b -> (a, p)
fBack b
db
                (c
dc, p
dpG) = d -> (c, p)
gBack d
dd
             in ((a
da, c
dc), (p, p) -> p
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
CB.plus (p
dpF, p
dpG))
        )
  {-# INLINE tensor #-}

instance (MergeZero (->) p) => Action (,) (DiffP p) where
  braid :: forall a b. DiffP p (a, b) (b, a)
braid = (p -> (a, b) -> ((b, a), (b, a) -> ((a, b), p)))
-> DiffP p (a, b) (b, a)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> (a, b) -> ((b, a), (b, a) -> ((a, b), p)))
 -> DiffP p (a, b) (b, a))
-> (p -> (a, b) -> ((b, a), (b, a) -> ((a, b), p)))
-> DiffP p (a, b) (b, a)
forall a b. (a -> b) -> a -> b
$ \p
_ (a
a, b
b) -> ((b
b, a
a), \(b
db, a
da) -> ((a
da, b
db), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE braid #-}

----------------------------------------------------------------------
-- Bimonoid structure
----------------------------------------------------------------------

-- | Copy in 'DiffP': forward copy, backward add.  The parameter gradient is
-- zero because copy has no parameters.
--
-- >>> let (_, pb) = runDiffP copy () (5 :: Int)
-- >>> pb (1, 2)
-- (3,())
instance (Merge (->) a, Zero (->) p) => Copy (DiffP p) a where
  copy :: DiffP p a (a, a)
copy = (p -> a -> ((a, a), (a, a) -> (a, p))) -> DiffP p a (a, a)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> ((a, a), (a, a) -> (a, p))) -> DiffP p a (a, a))
-> (p -> a -> ((a, a), (a, a) -> (a, p))) -> DiffP p a (a, a)
forall a b. (a -> b) -> a -> b
$ \p
_ a
a -> ((a
a, a
a), \(a
da1, a
da2) -> ((a, a) -> a
forall (arr :: * -> * -> *) a. Merge arr a => arr (a, a) a
CB.plus (a
da1, a
da2), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE copy #-}

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

-- | Add in 'DiffP': forward add, backward copy.  The parameter gradient is
-- zero because addition has no parameters.
--
-- >>> let (_, pb) = runDiffP plus () ((3, 4) :: (Int, Int))
-- >>> pb 1
-- ((1,1),())
instance (Merge (->) a, Zero (->) p) => Merge (DiffP p) a where
  plus :: DiffP p (a, a) a
plus = (p -> (a, a) -> (a, a -> ((a, a), p))) -> DiffP p (a, a) a
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> (a, a) -> (a, a -> ((a, a), p))) -> DiffP p (a, a) a)
-> (p -> (a, a) -> (a, a -> ((a, a), p))) -> DiffP p (a, a) a
forall a b. (a -> b) -> a -> b
$ \p
_ (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), () -> p
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ()))
  {-# INLINE plus #-}

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

----------------------------------------------------------------------
-- Wiring helpers
----------------------------------------------------------------------

-- | Add a residual connection around an operation.
--
--   forward:  y = x + op(x)
--   backward: dx = dy + dOp
residual :: (Num a) => DiffP p a a -> DiffP p a a
residual :: forall a p. Num a => DiffP p a a -> DiffP p a a
residual DiffP p a a
op = (p -> a -> (a, a -> (a, p))) -> DiffP p a a
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((p -> a -> (a, a -> (a, p))) -> DiffP p a a)
-> (p -> a -> (a, a -> (a, p))) -> DiffP p a a
forall a b. (a -> b) -> a -> b
$ \p
p a
a ->
  let (a
b, a -> (a, p)
opBack) = DiffP p a a -> p -> a -> (a, a -> (a, p))
forall p a b. DiffP p a b -> p -> a -> (b, b -> (a, p))
runDiffP DiffP p a a
op p
p a
a
   in ( a
a a -> a -> a
forall a. Num a => a -> a -> a
+ a
b,
        \a
dy ->
          let (a
daOp, p
dp) = a -> (a, p)
opBack a
dy
           in (a
dy a -> a -> a
forall a. Num a => a -> a -> a
+ a
daOp, p
dp)
      )
{-# INLINE residual #-}

-- | Split a parameter tuple so the left and right halves can be used by
-- independent parallel branches.
--
-- This is the parameter-product combinator that the fixed-parameter
-- 'Category' instance cannot express.  It keeps the migration semantics
-- identical to the original circuits-llm 'DiffP'.
splitP :: DiffP p1 a1 b1 -> DiffP p2 a2 b2 -> DiffP (p1, p2) (a1, a2) (b1, b2)
splitP :: forall p1 a1 b1 p2 a2 b2.
DiffP p1 a1 b1
-> DiffP p2 a2 b2 -> DiffP (p1, p2) (a1, a2) (b1, b2)
splitP (DiffP p1 -> a1 -> (b1, b1 -> (a1, p1))
f1) (DiffP p2 -> a2 -> (b2, b2 -> (a2, p2))
f2) = ((p1, p2)
 -> (a1, a2) -> ((b1, b2), (b1, b2) -> ((a1, a2), (p1, p2))))
-> DiffP (p1, p2) (a1, a2) (b1, b2)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP (((p1, p2)
  -> (a1, a2) -> ((b1, b2), (b1, b2) -> ((a1, a2), (p1, p2))))
 -> DiffP (p1, p2) (a1, a2) (b1, b2))
-> ((p1, p2)
    -> (a1, a2) -> ((b1, b2), (b1, b2) -> ((a1, a2), (p1, p2))))
-> DiffP (p1, p2) (a1, a2) (b1, b2)
forall a b. (a -> b) -> a -> b
$ \(p1
p1, p2
p2) (a1
a1, a2
a2) ->
  let (b1
b1, b1 -> (a1, p1)
back1) = p1 -> a1 -> (b1, b1 -> (a1, p1))
f1 p1
p1 a1
a1
      (b2
b2, b2 -> (a2, p2)
back2) = p2 -> a2 -> (b2, b2 -> (a2, p2))
f2 p2
p2 a2
a2
   in ( (b1
b1, b2
b2),
        \(b1
db1, b2
db2) ->
          let (a1
da1, p1
dp1) = b1 -> (a1, p1)
back1 b1
db1
              (a2
da2, p2
dp2) = b2 -> (a2, p2)
back2 b2
db2
           in ((a1
da1, a2
da2), (p1
dp1, p2
dp2))
      )
{-# INLINE splitP #-}

-- | Pair two operations that share the same input type but produce
-- independent outputs.  Input cotangents are added; parameter gradients are
-- paired.
joinP :: (Num a) => DiffP p1 a b1 -> DiffP p2 a b2 -> DiffP (p1, p2) a (b1, b2)
joinP :: forall a p1 b1 p2 b2.
Num a =>
DiffP p1 a b1 -> DiffP p2 a b2 -> DiffP (p1, p2) a (b1, b2)
joinP (DiffP p1 -> a -> (b1, b1 -> (a, p1))
f1) (DiffP p2 -> a -> (b2, b2 -> (a, p2))
f2) = ((p1, p2) -> a -> ((b1, b2), (b1, b2) -> (a, (p1, p2))))
-> DiffP (p1, p2) a (b1, b2)
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP (((p1, p2) -> a -> ((b1, b2), (b1, b2) -> (a, (p1, p2))))
 -> DiffP (p1, p2) a (b1, b2))
-> ((p1, p2) -> a -> ((b1, b2), (b1, b2) -> (a, (p1, p2))))
-> DiffP (p1, p2) a (b1, b2)
forall a b. (a -> b) -> a -> b
$ \(p1
p1, p2
p2) a
a ->
  let (b1
b1, b1 -> (a, p1)
back1) = p1 -> a -> (b1, b1 -> (a, p1))
f1 p1
p1 a
a
      (b2
b2, b2 -> (a, p2)
back2) = p2 -> a -> (b2, b2 -> (a, p2))
f2 p2
p2 a
a
   in ( (b1
b1, b2
b2),
        \(b1
db1, b2
db2) ->
          let (a
da1, p1
dp1) = b1 -> (a, p1)
back1 b1
db1
              (a
da2, p2
dp2) = b2 -> (a, p2)
back2 b2
db2
           in (a
da1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
da2, (p1
dp1, p2
dp2))
      )
{-# INLINE joinP #-}

----------------------------------------------------------------------
-- Relationship to the phantom-tagged Diff arrow
----------------------------------------------------------------------

-- | Embed a phantom-tagged 'Diff' into 'DiffP ()'.  The phantom tag is
-- discarded because 'DiffP' carries parameters at the value level.
--
-- >>> let d = Diff (\x -> (x * x, \dy -> 2 * x * dy)) :: Diff () Double Double
-- >>> let (y, pb) = runDiffP (toParam d) () 3.0
-- >>> y
-- 9.0
-- >>> pb 1.0
-- (6.0,())
toParam :: Diff q a b -> DiffP () a b
toParam :: forall {k} (q :: k) a b. Diff q a b -> DiffP () a b
toParam Diff q a b
d = (() -> a -> (b, b -> (a, ()))) -> DiffP () a b
forall p a b. (p -> a -> (b, b -> (a, p))) -> DiffP p a b
DiffP ((() -> a -> (b, b -> (a, ()))) -> DiffP () a b)
-> (() -> a -> (b, b -> (a, ()))) -> DiffP () a b
forall a b. (a -> b) -> a -> b
$ \()
_ a
a ->
  let (b
b, b -> a
pb) = Diff q a b -> a -> (b, b -> a)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff q a b
d a
a
   in (b
b, \b
db -> (b -> a
pb b
db, ()))
{-# INLINE toParam #-}

-- | Project a parameter-free 'DiffP ()' back into the phantom-tagged
-- 'Diff' arrow.
--
-- This is a section/retract pair with 'toParam':
-- @fromParam . toParam = id@ for the parameter-free fragment.
--
-- >>> let d = Diff (\x -> (x * x, \dy -> 2 * x * dy)) :: Diff () Double Double
-- >>> let (y, pb) = runDiff (fromParam (toParam d)) 3.0
-- >>> y
-- 9.0
-- >>> pb 1.0
-- 6.0
fromParam :: DiffP () a b -> Diff () a b
fromParam :: forall a b. DiffP () a b -> Diff () a b
fromParam (DiffP () -> a -> (b, b -> (a, ()))
f) = (a -> (b, b -> a)) -> Diff () a b
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((a -> (b, b -> a)) -> Diff () a b)
-> (a -> (b, b -> a)) -> Diff () a b
forall a b. (a -> b) -> a -> b
$ \a
a ->
  let (b
b, b -> (a, ())
pb) = () -> a -> (b, b -> (a, ()))
f () a
a
   in (b
b, (a, ()) -> a
forall a b. (a, b) -> a
fst ((a, ()) -> a) -> (b -> (a, ())) -> b -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. b -> (a, ())
pb)
{-# INLINE fromParam #-}