Explicit.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_HADDOCK not-home #-}

-- |
-- Module      : Numeric.Backprop.Explicit
-- Copyright   : (c) Justin Le 2023
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Provides "explicit" versions of all of the functions in
-- "Numeric.Backprop".  Instead of relying on a 'Backprop' instance, allows
-- you to manually provide 'zero', 'add', and 'one' on a per-value basis.
--
-- It is recommended you use "Numeric.Backprop" or "Numeric.Backprop.Num"
-- instead, unless your type has no 'Num' instance, or you else you want to
-- avoid defining orphan 'Backprop' instances for external types.  Can also
-- be useful if mixing and matching styles.
--
-- See "Numeric.Backprop" for fuller documentation on using these
-- functions.
--
-- WARNING: API of this module can be considered only "semi-stable"; while
-- the API of "Numeric.Backprop" and "Numeric.Backprop.Num" are kept
-- consistent, some argument order changes might happen in this module to
-- reflect changes in underlying implementation.
--
-- @since 0.2.0.0
module Numeric.Backprop.Explicit (
  -- * Types
  BVar,
  W,
  Backprop (..),
  ABP (..),
  NumBP (..),

  -- * Explicit 'zero', 'add', and 'one'
  ZeroFunc (..),
  zfNum,
  zfNums,
  zeroFunc,
  zeroFuncs,
  zfFunctor,
  AddFunc (..),
  afNum,
  afNums,
  addFunc,
  addFuncs,
  OneFunc (..),
  ofNum,
  ofNums,
  oneFunc,
  oneFuncs,
  ofFunctor,

  -- * Running
  backprop,
  evalBP,
  gradBP,
  backpropWith,

  -- ** Multiple inputs
  evalBP0,
  backprop2,
  evalBP2,
  gradBP2,
  backpropWith2,
  backpropN,
  evalBPN,
  gradBPN,
  backpropWithN,
  RPureConstrained,

  -- * Manipulating 'BVar'
  constVar,
  auto,
  coerceVar,
  viewVar,
  setVar,
  overVar,
  sequenceVar,
  collectVar,
  previewVar,
  toListOfVar,

  -- ** With Isomorphisms
  isoVar,
  isoVar2,
  isoVar3,
  isoVarN,

  -- ** With 'Op's
  liftOp,
  liftOp1,
  liftOp2,
  liftOp3,

  -- ** Generics
  splitBV,
  joinBV,
  BVGroup,

  -- * 'Op'
  Op (..),

  -- ** Creation
  op0,
  opConst,
  idOp,
  bpOp,

  -- *** Giving gradients directly
  op1,
  op2,
  op3,

  -- *** From Isomorphisms
  opCoerce,
  opTup,
  opIso,
  opIsoN,
  opLens,

  -- *** No gradients
  noGrad1,
  noGrad,

  -- * Utility
  Reifies,
) where

import Data.Bifunctor
import Data.Functor.Identity
import Data.Reflection
import Data.Type.Util
import Data.Vinyl.Core
import Data.Vinyl.TypeLevel
import GHC.Generics as G
import Lens.Micro
import Numeric.Backprop.Class
import Numeric.Backprop.Internal
import Numeric.Backprop.Op
import Unsafe.Coerce

-- | 'ZeroFunc's for every item in a type level list based on their
-- 'Num' instances
--
-- @since 0.2.0.0
zfNums :: RPureConstrained Num as => Rec ZeroFunc as
zfNums = rpureConstrained @Num zfNum

-- | 'zeroFunc' for instances of 'Functor'
--
-- @since 0.2.1.0
zfFunctor :: (Backprop a, Functor f) => ZeroFunc (f a)
zfFunctor = ZF zeroFunctor
{-# INLINE zfFunctor #-}

-- | 'ZeroFunc's for every item in a type level list based on their
-- 'Num' instances
--
-- @since 0.2.0.0
afNums :: RPureConstrained Num as => Rec AddFunc as
afNums = rpureConstrained @Num afNum

-- | 'ZeroFunc's for every item in a type level list based on their
-- 'Num' instances
--
-- @since 0.2.0.0
ofNums :: RPureConstrained Num as => Rec OneFunc as
ofNums = rpureConstrained @Num ofNum

-- | 'OneFunc' for instances of 'Functor'
--
-- @since 0.2.1.0
ofFunctor :: (Backprop a, Functor f) => OneFunc (f a)
ofFunctor = OF oneFunctor
{-# INLINE ofFunctor #-}

-- | Generate an 'ZeroFunc' for every type in a type-level list, if every
-- type has an instance of 'Backprop'.
--
-- @since 0.2.0.0
zeroFuncs :: RPureConstrained Backprop as => Rec ZeroFunc as
zeroFuncs = rpureConstrained @Backprop zeroFunc

-- | Generate an 'AddFunc' for every type in a type-level list, if every
-- type has an instance of 'Backprop'.
--
-- @since 0.2.0.0
addFuncs :: RPureConstrained Backprop as => Rec AddFunc as
addFuncs = rpureConstrained @Backprop addFunc

-- | Generate an 'OneFunc' for every type in a type-level list, if every
-- type has an instance of 'Backprop'.
--
-- @since 0.2.0.0
oneFuncs :: RPureConstrained Backprop as => Rec OneFunc as
oneFuncs = rpureConstrained @Backprop oneFunc

-- | Shorter alias for 'constVar', inspired by the /ad/ library.
--
-- @since 0.2.0.0
auto :: a -> BVar s a
auto = constVar
{-# INLINE auto #-}

-- | 'Numeric.Backprop.backpropN', but with explicit 'zero' and 'one'.
backpropN ::
  forall as b.
  () =>
  Rec ZeroFunc as ->
  OneFunc b ->
  (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) ->
  Rec Identity as ->
  (b, Rec Identity as)
backpropN zfs ob f xs = case backpropWithN zfs f xs of
  (y, g) -> (y, g (runOF ob y))
{-# INLINE backpropN #-}

-- | 'Numeric.Backprop.backprop', but with explicit 'zero' and 'one'.
backprop ::
  ZeroFunc a ->
  OneFunc b ->
  (forall s. Reifies s W => BVar s a -> BVar s b) ->
  a ->
  (b, a)
backprop zfa ofb f =
  second (\case Identity x :& RNil -> x)
    . backpropN (zfa :& RNil) ofb (f . (\case x :& RNil -> x))
    . (:& RNil)
    . Identity
{-# INLINE backprop #-}

-- | 'Numeric.Backprop.backpropWith', but with explicit 'zero'.
--
-- Note that argument order changed in v0.2.4.
backpropWith ::
  ZeroFunc a ->
  (forall s. Reifies s W => BVar s a -> BVar s b) ->
  a ->
  (b, b -> a)
backpropWith zfa f =
  second ((\case Identity x :& RNil -> x) .)
    . backpropWithN (zfa :& RNil) (f . (\case x :& RNil -> x))
    . (:& RNil)
    . Identity
{-# INLINE backpropWith #-}

-- | 'evalBP' but with no arguments.  Useful when everything is just given
-- through 'constVar'.
evalBP0 :: (forall s. Reifies s W => BVar s a) -> a
evalBP0 x = evalBPN (const x) RNil
{-# INLINE evalBP0 #-}

-- | Turn a function @'BVar' s a -> 'BVar' s b@ into the function @a -> b@
-- that it represents.
--
-- Benchmarks show that this should have virtually no overhead over
-- directly writing a @a -> b@. 'BVar' is, in this situation, a zero-cost
-- abstraction, performance-wise.
--
-- See documentation of 'Numeric.Backprop.backprop' for more information.
evalBP :: (forall s. Reifies s W => BVar s a -> BVar s b) -> a -> b
evalBP f = evalBPN (f . (\case x :& RNil -> x)) . (:& RNil) . Identity
{-# INLINE evalBP #-}

-- | 'Numeric.Backprop.gradBP', but with explicit 'zero' and 'one'.
gradBP ::
  ZeroFunc a ->
  OneFunc b ->
  (forall s. Reifies s W => BVar s a -> BVar s b) ->
  a ->
  a
gradBP zfa ofb f = snd . backprop zfa ofb f
{-# INLINE gradBP #-}

-- | 'Numeric.Backprop.gradBP', Nbut with explicit 'zero' and 'one'.
gradBPN ::
  Rec ZeroFunc as ->
  OneFunc b ->
  (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) ->
  Rec Identity as ->
  Rec Identity as
gradBPN zfas ofb f = snd . backpropN zfas ofb f
{-# INLINE gradBPN #-}

-- | 'Numeric.Backprop.backprop2', but with explicit 'zero' and 'one'.
backprop2 ::
  ZeroFunc a ->
  ZeroFunc b ->
  OneFunc c ->
  (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
  a ->
  b ->
  (c, (a, b))
backprop2 zfa zfb ofc f x y =
  second (\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) $
    backpropN
      (zfa :& zfb :& RNil)
      ofc
      (\(x' :& y' :& RNil) -> f x' y')
      (Identity x :& Identity y :& RNil)
{-# INLINE backprop2 #-}

-- | 'Numeric.Backprop.backpropWith2', but with explicit 'zero'.
--
-- Note that argument order changed in v0.2.4.
--
-- @since 0.2.0.0
backpropWith2 ::
  ZeroFunc a ->
  ZeroFunc b ->
  (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
  a ->
  b ->
  (c, c -> (a, b))
backpropWith2 zfa zfb f x y =
  second ((\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) .) $
    backpropWithN
      (zfa :& zfb :& RNil)
      (\(x' :& y' :& RNil) -> f x' y')
      (Identity x :& Identity y :& RNil)
{-# INLINE backpropWith2 #-}

-- | 'evalBP' for a two-argument function.  See
-- 'Numeric.Backprop.backprop2' for notes.
evalBP2 ::
  (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
  a ->
  b ->
  c
evalBP2 f x y =
  evalBPN (\(x' :& y' :& RNil) -> f x' y') $
    Identity x
      :& Identity y
      :& RNil
{-# INLINE evalBP2 #-}

-- | 'Numeric.Backprop.gradBP2' with explicit 'zero' and 'one'.
gradBP2 ::
  ZeroFunc a ->
  ZeroFunc b ->
  OneFunc c ->
  (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
  a ->
  b ->
  (a, b)
gradBP2 zfa zfb ofc f x = snd . backprop2 zfa zfb ofc f x
{-# INLINE gradBP2 #-}

-- | 'Numeric.Backprop.bpOp' with explicit 'zero'.
bpOp ::
  Rec ZeroFunc as ->
  (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) ->
  Op as b
bpOp zfs f = Op (backpropWithN zfs f)
{-# INLINE bpOp #-}

-- | 'Numeric.Backprop.overVar' with explicit 'add' and 'zero'.
--
-- @since 0.2.4.0
overVar ::
  Reifies s W =>
  AddFunc a ->
  AddFunc b ->
  ZeroFunc a ->
  ZeroFunc b ->
  Lens' b a ->
  (BVar s a -> BVar s a) ->
  BVar s b ->
  BVar s b
overVar afa afb zfa zfb l f x = setVar afa afb zfa l (f (viewVar afa zfb l x)) x
{-# INLINE overVar #-}

-- | 'Numeric.Backprop.isoVar' with explicit 'add' and 'zero'.
isoVar ::
  Reifies s W =>
  AddFunc a ->
  (a -> b) ->
  (b -> a) ->
  BVar s a ->
  BVar s b
isoVar af f g = liftOp1 af (opIso f g)
{-# INLINE isoVar #-}

-- | 'Numeric.Backprop.isoVar2' with explicit 'add' and 'zero'.
isoVar2 ::
  Reifies s W =>
  AddFunc a ->
  AddFunc b ->
  (a -> b -> c) ->
  (c -> (a, b)) ->
  BVar s a ->
  BVar s b ->
  BVar s c
isoVar2 afa afb f g = liftOp2 afa afb (opIso2 f g)
{-# INLINE isoVar2 #-}

-- | 'Numeric.Backprop.isoVar3' with explicit 'add' and 'zero'.
isoVar3 ::
  Reifies s W =>
  AddFunc a ->
  AddFunc b ->
  AddFunc c ->
  (a -> b -> c -> d) ->
  (d -> (a, b, c)) ->
  BVar s a ->
  BVar s b ->
  BVar s c ->
  BVar s d
isoVar3 afa afb afc f g = liftOp3 afa afb afc (opIso3 f g)
{-# INLINE isoVar3 #-}

-- | 'Numeric.Backprop.isoVarN' with explicit 'add' and 'zero'.
isoVarN ::
  Reifies s W =>
  Rec AddFunc as ->
  (Rec Identity as -> b) ->
  (b -> Rec Identity as) ->
  Rec (BVar s) as ->
  BVar s b
isoVarN afs f g = liftOp afs (opIsoN f g)
{-# INLINE isoVarN #-}

-- | Helper class for generically "splitting" and "joining" 'BVar's into
-- constructors.  See 'Numeric.Backprop.splitBV' and
-- 'Numeric.Backprop.joinBV'.
--
-- See "Numeric.Backprop#hkd" for a tutorial on how to use this.
--
-- Instances should be available for types made with one constructor whose
-- fields are all instances of 'Backprop', with a 'Generic' instance.
--
-- @since 0.2.2.0
class BVGroup s as i o | o -> i, i -> as where
  -- | Helper method for generically "splitting" 'BVar's out of
  -- constructors inside a 'BVar'.  See 'splitBV'.
  gsplitBV :: Rec AddFunc as -> Rec ZeroFunc as -> BVar s (i ()) -> o ()

  -- | Helper method for generically "joining" 'BVar's inside
  -- a constructor into a 'BVar'.  See 'joinBV'.
  gjoinBV :: Rec AddFunc as -> Rec ZeroFunc as -> o () -> BVar s (i ())

instance BVGroup s '[] (K1 i a) (K1 i (BVar s a)) where
  gsplitBV _ _ = K1 . coerceVar
  {-# INLINE gsplitBV #-}
  gjoinBV _ _ = coerceVar . unK1
  {-# INLINE gjoinBV #-}

instance
  BVGroup s as i o =>
  BVGroup s as (M1 p c i) (M1 p c o)
  where
  gsplitBV afs zfs = M1 . gsplitBV afs zfs . coerceVar @_ @(i ())
  {-# INLINE gsplitBV #-}
  gjoinBV afs zfs = coerceVar @(i ()) . gjoinBV afs zfs . unM1
  {-# INLINE gjoinBV #-}

instance BVGroup s '[] V1 V1 where
  gsplitBV _ _ = unsafeCoerce
  {-# INLINE gsplitBV #-}
  gjoinBV _ _ = \case {}
  {-# INLINE gjoinBV #-}

instance BVGroup s '[] U1 U1 where
  gsplitBV _ _ _ = U1
  {-# INLINE gsplitBV #-}
  gjoinBV _ _ _ = constVar U1
  {-# INLINE gjoinBV #-}

instance
  ( Reifies s W
  , BVGroup s as i1 o1
  , BVGroup s bs i2 o2
  , cs ~ (as ++ bs)
  , RecApplicative as
  ) =>
  BVGroup s (i1 () ': i2 () ': cs) (i1 :*: i2) (o1 :*: o2)
  where
  gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy = x :*: y
    where
      (afas, afbs) = splitRec afs
      (zfas, zfbs) = splitRec zfs
      zfab = ZF $ \(xx :*: yy) -> runZF zfa xx :*: runZF zfb yy
      x = gsplitBV afas zfas . viewVar afa zfab p1 $ xy
      y = gsplitBV afbs zfbs . viewVar afb zfab p2 $ xy
  {-# INLINE gsplitBV #-}
  gjoinBV (afa :& afb :& afs) (_ :& _ :& zfs) (x :*: y) =
    isoVar2
      afa
      afb
      (:*:)
      unP
      (gjoinBV afas zfas x)
      (gjoinBV afbs zfbs y)
    where
      (afas, afbs) = splitRec afs
      (zfas, zfbs) = splitRec zfs
      unP (xx :*: yy) = (xx, yy)
  {-# INLINE gjoinBV #-}

-- | This instance is possible but it is not clear when it would be useful
instance
  ( Reifies s W
  , BVGroup s as i1 o1
  , BVGroup s bs i2 o2
  , cs ~ (as ++ bs)
  , RecApplicative as
  ) =>
  BVGroup s (i1 () ': i2 () ': cs) (i1 :+: i2) (o1 :+: o2)
  where
  gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy =
    case previewVar afa zf s1 xy of
      Just x -> L1 $ gsplitBV afas zfas x
      Nothing -> case previewVar afb zf s2 xy of
        Just y -> R1 $ gsplitBV afbs zfbs y
        Nothing -> error "Numeric.Backprop.gsplitBV: Internal error occurred"
    where
      zf = ZF $ \case
        L1 xx -> L1 $ runZF zfa xx
        R1 yy -> R1 $ runZF zfb yy
      (afas, afbs) = splitRec afs
      (zfas, zfbs) = splitRec zfs
  {-# INLINE gsplitBV #-}
  gjoinBV (afa :& afb :& afs) (zfa :& zfb :& zfs) = \case
    L1 x ->
      liftOp1
        afa
        (op1 (\xx -> (L1 xx, \case L1 d -> d; R1 _ -> runZF zfa xx)))
        (gjoinBV afas zfas x)
    R1 y ->
      liftOp1
        afb
        (op1 (\yy -> (R1 yy, \case L1 _ -> runZF zfb yy; R1 d -> d)))
        (gjoinBV afbs zfbs y)
    where
      (afas, afbs) = splitRec afs
      (zfas, zfbs) = splitRec zfs
  {-# INLINE gjoinBV #-}

-- | 'Numeric.Backprop.splitBV' with explicit 'add' and 'zero'.
--
-- @since 0.2.2.0
splitBV ::
  forall z f s as.
  ( Generic (z f)
  , Generic (z (BVar s))
  , BVGroup s as (Rep (z f)) (Rep (z (BVar s)))
  , Reifies s W
  ) =>
  AddFunc (Rep (z f) ()) ->
  Rec AddFunc as ->
  ZeroFunc (z f) ->
  Rec ZeroFunc as ->
  -- | 'BVar' of value
  BVar s (z f) ->
  -- | 'BVar's of fields
  z (BVar s)
splitBV af afs zf zfs =
  G.to
    . gsplitBV afs zfs
    . viewVar af zf (lens (from @(z f) @()) (const G.to))
{-# INLINE splitBV #-}

-- | 'Numeric.Backprop.joinBV' with explicit 'add' and 'zero'.
--
-- @since 0.2.2.0
joinBV ::
  forall z f s as.
  ( Generic (z f)
  , Generic (z (BVar s))
  , BVGroup s as (Rep (z f)) (Rep (z (BVar s)))
  , Reifies s W
  ) =>
  AddFunc (z f) ->
  Rec AddFunc as ->
  ZeroFunc (Rep (z f) ()) ->
  Rec ZeroFunc as ->
  -- | 'BVar's of fields
  z (BVar s) ->
  -- | 'BVar' of combined value
  BVar s (z f)
joinBV af afs zf zfs =
  viewVar af zf (lens G.to (const from))
    . gjoinBV afs zfs
    . from @(z (BVar s)) @()
{-# INLINE joinBV #-}