Generality

[gdalle]

• Should we support primitives beyond pushforwards / pullbacks ?
  • Requires making assumptions on argument types
• Should we support higher order derivatives?
  • Requires mixing backends, ergo a non-empty interface
• Should we support argument types beyond arrays / numbers?
  • Makes it harder to reason about mutability and operators
• Should we support multiple arguments?
  • Caused a lot of trouble in AbstractDifferentiation
  • Some backends don't like this anyway (ForwardDiff, ReverseDiff)
• What should be the function signature?
  • f(x::Array) -> y::Array as right now (more or less)
  • f!(y::Array, x::Array) -> z::Nothing as in ForwardDiff
  • f!(y::Array, x::Array) -> z::Union{Number,Nothing} as in Enzyme

[adrhill]

(1) Mutating pushforwards & pullbacks

These are at the core of DifferentiationInterface. Each backend needs to implement either a pushforward!, pullback! or both.

(1.1) Mutation of tangent

pushforward!(dy, backend, f, x, dx) -> (y, dy)
pullback!(dx, backend, f, x, dy) -> (y, dx)

• Q: Do we want to support multiple arguments?
• Q: If so, how? Tuples? (See #53)
• Q: What do we understand by mutation? Does this simply provide a "scratch space" in memory for the function or does the result have to be written there?
• Q: Does the pushforward abstraction fully cover ForwardDiff and FiniteDiff/erences?

(1.2) Mutation of primal and tangent

pushforward!(y, dy, backend, f!, x, dx) -> (y, dy)
pullback!(y, dx, backend, f!, x, dy) -> (y, dx)

(2) Allocating tangents

These are "obtained for free" using (1) as a fallback.

pushforward(backend, f, x, dx) -> (y, dy)
pullback(backend, f, x, dy) -> (y, dx)

The call order can be reversed ((1) calling (2)) in case mutation is not supported.

(3) Pullback closures

These look like a duplicate of (1), but are central to reverse-mode AD systems like Zygote to avoid redundant forward passes. However, they are seemingly irrelevant to Enzyme, which DifferentiationInterface wants to fully support.

pullback_function(backend, f, x, dy) -> (y, f_pb)

• Q: Do we completely ignore this?
• Q: Do we provide different methods for these pullback closures, e.g. a mutating and non-mutating one?

(4) Jacobian

Jacobians are "obtained for free" using (1) and (2) as fallbacks.
For $f: \mathbb{R}^n \rightarrow \mathbb{R}^m$ the corresponding $m\times n$ Jacobian can be constructed column-wise via $n$ calls to pushforward or row-wise via $m$ calls to pullback

• implement value_and_jacobian and jacobian via multiple dispatch on ForwardBackend (pushforward) or ReverseBackend (pullback)
• Q: Should we apply a heuristic based on dimensions $n$ and $m$ for backends that support both forward and reverse? Or should we make the user choose by separating the backends into strict forward and reverse, e.g. EnzymeForwardBackend and EnzymeReverseBackend?
• implement in-place versions value_and_jacobian! and jacobian! via fallback to pushforward! and pullback!
• Q: For some backends like Zygote, (3) is a much better fallback than (1) and (2), avoiding redundant forward-passes. Should (1) still be the default fallback?
• Q: Is it possible to use a unified type of standard basis vectors across backends?

(5) Gradient

Gradients are a special case of the Jacobian for $m=1$ and require $n$ calls to pushforward or one call to pullback

• implement value_and_gradient and gradient via multiple dispatch on ForwardBackend (pushforward) or ReverseBackend (pullback)
• implement in-place versions value_and_gradient! and gradient! via fallback to pushforward! and pullback!

[adrhill]

Regarding (3), it looks like Enzyme supports "reverse pass punctions" via ReverseModeSplit and autodiff_deferred_thunk / autodiff_thunk.

CC @wsmoses

[adrhill]

I think this can largely be considered closed by #23 and should be replaced by more specific issues.