-- | Predefined common functions for simplification and pretty-printing of AST.
module HordeAd.AstEngine
  ( -- * The joint inlining and simplification term transformations
    simplifyArtifact, simplifyArtifactGradient, simplifyArtifactDerivative
  , simplifyInline, simplifyInlineContract, simplifyInlineContractNoExpand
    -- * Pretty-printing terms in a few useful configurations
  , printAstVarName
  , printAstSimple, printAstPretty, printAstPrettyButNested
  , printArtifactSimple, printArtifactPretty
  , printArtifactPrimalSimple, printArtifactPrimalPretty
  ) where

import Prelude

import Data.EnumMap.Strict qualified as EM
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IM

import HordeAd.Core.Ast
import HordeAd.Core.AstInline
import HordeAd.Core.AstPrettyPrint
import HordeAd.Core.AstTraverse

-- * The joint inlining and simplification term transformation

-- | Simplify the whole reverse derivative artifact (which includes
-- also the primal value computed during the differentiation process).
simplifyArtifact :: forall x z.
                    AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev x z
art =
  let !der :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
der = AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract (AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
 -> AstTensor AstMethodLet PrimalSpan (ADTensorKind x))
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artDerivativeRev AstArtifactRev x z
art in
  let !prim :: AstTensor AstMethodLet PrimalSpan z
prim = AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract (AstTensor AstMethodLet PrimalSpan z
 -> AstTensor AstMethodLet PrimalSpan z)
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall a b. (a -> b) -> a -> b
$ AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artPrimalRev AstArtifactRev x z
art
  in AstArtifactRev x z
art {artDerivativeRev = der, artPrimalRev = prim}

-- | Simplify only the gradient in the reverse derivative artifact.
simplifyArtifactGradient :: forall x z.
                            AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifactGradient :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifactGradient AstArtifactRev x z
art =
  AstArtifactRev x z
art { artDerivativeRev =
        simplifyInlineContract $ artDerivativeRev art }

-- | Simplify only the derivative in the forward derivative artifact.
simplifyArtifactDerivative :: forall x z.
                              AstArtifactFwd x z -> AstArtifactFwd x z
simplifyArtifactDerivative :: forall (x :: TK) (z :: TK).
AstArtifactFwd x z -> AstArtifactFwd x z
simplifyArtifactDerivative AstArtifactFwd x z
art =
  AstArtifactFwd x z
art { artDerivativeFwd =
        simplifyInlineContract $ artDerivativeFwd art }

-- | A mixture of simplification and inlining to use when the resultng
-- term is not yet supposed to be interpreted using a computational backed,
-- but rather to be stored and later composed with other terms.
simplifyInline
  :: forall z s. AstSpan s
  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline :: forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline =
  AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstMemo, AstTensor AstMethodLet s z) -> AstTensor AstMethodLet s z
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z
    -> (AstMemo, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty
  (AstTensor AstMethodLet s z
 -> (AstMemo, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstMemo, AstTensor AstMethodLet s z) -> AstTensor AstMethodLet s z
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z
    -> (AstMemo, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty
  (AstTensor AstMethodLet s z
 -> (AstMemo, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst

-- | A mixture of simplification, inlining and recognition of additional
-- backend-specific primitives, to be used just before a term
-- is interpreted as a value in the computational backend.
simplifyInlineContract
  :: forall z s. AstSpan s
  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract :: forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract =
  AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstMemo, AstTensor AstMethodLet s z) -> AstTensor AstMethodLet s z
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z
    -> (AstMemo, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty
  (AstTensor AstMethodLet s z
 -> (AstMemo, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstMemo, AstTensor AstMethodLet s z) -> AstTensor AstMethodLet s z
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z
    -> (AstMemo, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty
  (AstTensor AstMethodLet s z
 -> (AstMemo, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst

simplifyInlineContractNoExpand
  :: forall z s. AstSpan s
  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContractNoExpand :: forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContractNoExpand =
  AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstMemo, AstTensor AstMethodLet s z) -> AstTensor AstMethodLet s z
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z
    -> (AstMemo, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty
  (AstTensor AstMethodLet s z
 -> (AstMemo, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstMemo, AstTensor AstMethodLet s z) -> AstTensor AstMethodLet s z
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s z
    -> (AstMemo, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty
  (AstTensor AstMethodLet s z
 -> (AstMemo, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst


-- * Pretty-printing terms in a few useful configurations

printAstVarRename :: AstSpan s
                  => IntMap String -> AstVarName s y -> String
printAstVarRename :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
IntMap String -> AstVarName s y -> String
printAstVarRename IntMap String
renames AstVarName s y
var =
  PrintConfig -> AstVarName s y -> ShowS
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
PrintConfig -> AstVarName s y -> ShowS
printAstVar (PrintConfig
defaulPrintConfig {varRenames = renames}) AstVarName s y
var String
""

printAstSimpleRename :: AstSpan s
                     => IntMap String -> AstTensor ms s y -> String
printAstSimpleRename :: forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
IntMap String -> AstTensor ms s y -> String
printAstSimpleRename IntMap String
renames AstTensor ms s y
t =
  PrintConfig -> Int -> AstTensor ms s y -> ShowS
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
PrintConfig -> Int -> AstTensor ms s y -> ShowS
printAst
    (PrintConfig
defaulPrintConfig {loseRoudtrip = False, varRenames = renames}) Int
0 AstTensor ms s y
t String
""

printAstPrettyRename :: AstSpan s
                     => IntMap String -> AstTensor ms s y -> String
printAstPrettyRename :: forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
IntMap String -> AstTensor ms s y -> String
printAstPrettyRename IntMap String
renames AstTensor ms s y
t =
  PrintConfig -> Int -> AstTensor ms s y -> ShowS
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
PrintConfig -> Int -> AstTensor ms s y -> ShowS
printAst (PrintConfig
defaulPrintConfig {varRenames = renames}) Int
0 AstTensor ms s y
t String
""

printAstVarName :: AstSpan s
                => AstVarName s y -> String
printAstVarName :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstVarName s y -> String
printAstVarName AstVarName s y
var =
  PrintConfig -> AstVarName s y -> ShowS
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
PrintConfig -> AstVarName s y -> ShowS
printAstVar PrintConfig
defaulPrintConfig AstVarName s y
var String
""

-- | Print an AST term in a form close to being able to roundtrip,
-- including explicit sharing preservation.
printAstSimple :: AstSpan s
               => AstTensor ms s y -> String
printAstSimple :: forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstSimple AstTensor ms s y
t =
  PrintConfig -> Int -> AstTensor ms s y -> ShowS
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
PrintConfig -> Int -> AstTensor ms s y -> ShowS
printAst (PrintConfig
defaulPrintConfig {loseRoudtrip = False}) Int
0 AstTensor ms s y
t String
""

-- | Print an AST term in a readable form that does not roundtrip,
-- and where Haskell @let@ (sharing on Haskell heap) is used instead
-- of explicit sharing of subterms.
printAstPretty :: AstSpan s
               => AstTensor ms s y -> String
printAstPretty :: forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor ms s y
t =
  PrintConfig -> Int -> AstTensor ms s y -> ShowS
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
PrintConfig -> Int -> AstTensor ms s y -> ShowS
printAst PrintConfig
defaulPrintConfig Int
0 AstTensor ms s y
t String
""

printAstPrettyButNested :: AstSpan s
                        => AstTensor ms s y -> String
printAstPrettyButNested :: forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPrettyButNested AstTensor ms s y
t =
  PrintConfig -> Int -> AstTensor ms s y -> ShowS
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
PrintConfig -> Int -> AstTensor ms s y -> ShowS
printAst (PrintConfig
defaulPrintConfig {ignoreNestedLambdas = False}) Int
0 AstTensor ms s y
t String
""

printArtifactSimple :: AstArtifactRev x z -> String
printArtifactSimple :: forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactSimple AstArtifactRev{AstTensor AstMethodLet PrimalSpan z
AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind z)
artDerivativeRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artVarDtRev :: AstVarName PrimalSpan (ADTensorKind z)
artVarDomainRev :: AstVarName PrimalSpan x
artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: AstTensor AstMethodLet PrimalSpan z
artVarDomainRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan x
artVarDtRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan (ADTensorKind z)
..} =
  let nDt :: Int
nDt = AstVarId -> Int
forall a. Enum a => a -> Int
fromEnum (AstVarName PrimalSpan (ADTensorKind z) -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
100000000
      renames :: IntMap String
renames = Int -> String -> IntMap String
forall a. Int -> a -> IntMap a
IM.singleton Int
nDt String
"dret"
      varsPP :: [String]
varsPP = [ IntMap String -> AstVarName PrimalSpan (ADTensorKind z) -> String
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
IntMap String -> AstVarName s y -> String
printAstVarRename IntMap String
renames AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev
               , IntMap String -> AstVarName PrimalSpan x -> String
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
IntMap String -> AstVarName s y -> String
printAstVarRename IntMap String
renames AstVarName PrimalSpan x
artVarDomainRev ]
  in String
"\\" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords [String]
varsPP
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ IntMap String
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
IntMap String -> AstTensor ms s y -> String
printAstSimpleRename IntMap String
renames AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artDerivativeRev

printArtifactPretty :: AstArtifactRev x z -> String
printArtifactPretty :: forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev{AstTensor AstMethodLet PrimalSpan z
AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind z)
artDerivativeRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artVarDomainRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan x
artVarDtRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev :: AstVarName PrimalSpan (ADTensorKind z)
artVarDomainRev :: AstVarName PrimalSpan x
artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: AstTensor AstMethodLet PrimalSpan z
..} =
  let nDt :: Int
nDt = AstVarId -> Int
forall a. Enum a => a -> Int
fromEnum (AstVarName PrimalSpan (ADTensorKind z) -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
100000000
      renames :: IntMap String
renames = Int -> String -> IntMap String
forall a. Int -> a -> IntMap a
IM.singleton Int
nDt String
"dret"
      varsPP :: [String]
varsPP = [ IntMap String -> AstVarName PrimalSpan (ADTensorKind z) -> String
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
IntMap String -> AstVarName s y -> String
printAstVarRename IntMap String
renames AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev
               , IntMap String -> AstVarName PrimalSpan x -> String
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
IntMap String -> AstVarName s y -> String
printAstVarRename IntMap String
renames AstVarName PrimalSpan x
artVarDomainRev ]
  in String
"\\" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords [String]
varsPP
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ IntMap String
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
IntMap String -> AstTensor ms s y -> String
printAstPrettyRename IntMap String
renames AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artDerivativeRev

printArtifactPrimalSimple :: AstArtifactRev x z -> String
printArtifactPrimalSimple :: forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalSimple AstArtifactRev{AstTensor AstMethodLet PrimalSpan z
AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind z)
artDerivativeRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artVarDomainRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan x
artVarDtRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev :: AstVarName PrimalSpan (ADTensorKind z)
artVarDomainRev :: AstVarName PrimalSpan x
artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: AstTensor AstMethodLet PrimalSpan z
..} =
  String
"\\" String -> ShowS
forall a. [a] -> [a] -> [a]
++ AstVarName PrimalSpan x -> String
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstVarName s y -> String
printAstVarName AstVarName PrimalSpan x
artVarDomainRev
       String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ AstTensor AstMethodLet PrimalSpan z -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstSimple AstTensor AstMethodLet PrimalSpan z
artPrimalRev

printArtifactPrimalPretty :: AstArtifactRev x z -> String
printArtifactPrimalPretty :: forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev{AstTensor AstMethodLet PrimalSpan z
AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind z)
artDerivativeRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artVarDomainRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan x
artVarDtRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev :: AstVarName PrimalSpan (ADTensorKind z)
artVarDomainRev :: AstVarName PrimalSpan x
artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: AstTensor AstMethodLet PrimalSpan z
..} =
  String
"\\" String -> ShowS
forall a. [a] -> [a] -> [a]
++ AstVarName PrimalSpan x -> String
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstVarName s y -> String
printAstVarName AstVarName PrimalSpan x
artVarDomainRev
       String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ AstTensor AstMethodLet PrimalSpan z -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor AstMethodLet PrimalSpan z
artPrimalRev