Redesign adjoint / transpose interface and algorithms

[fritzo]

Design doc

Reading list

• https://popl21.sigplan.org/details/lafi-2021-papers/9/Decomposing-reverse-mode-automatic-differentiation also The Autodiff Cookbook
• The Simple Essence of Automatic Differentiation
• Semiring programming
• Two tricks for the price of one: Linear filters and their transposes

Questions

1. Linear Types: Funsor metadata includes domains for inputs and output, similar to a type signature a1 * ... * an -> b. Should metadata also indicate which inputs are treated linearly, similar to a type signature a1 * ... * an -> (b1 * ... * bn -o c)?
2. linearize(): How can we implement a linearize(fn: Funsor, linear_vars: Frozenset[Variable])? Can this introduce new patterns but avoid introducing new Funsor subclasses?
3. transpose(): How can we implement a transpose(fn:Funsor, ???)? Does this require new Funsor subclasses?
4. Types: What is the type of backward-marginals? Of backward-sample? Of backward-argmax? Of backward-kbest?
5. Decomposition: Can we decompose backward-sample into a Monte Carlo interpretation of backward-marginal? Can we decompose backward-argmax into a "maximize" interpretation of backward-marginal?
6. Can we implement a tape-free adjoint algorithm, as suggested in Elliot's "Simple Essence of AD"?

[ordabayevy]

Just throwing out an idea here. Another perspective on two different semi-rings (ops.add, ops.mul) and (ops.logaddexp, ops.add) is to view the latter as operations lifted by a Log Functor. In Haskell notation:

instance Functor Log where
    fmap f :: Log a -> Log b
    fmap f (Log a) = Log (f a)

...

fmap (*) (Log (x, y)) = Log (x*y) = Log x + Log y
fmap (+)  = logaddexp

These two different types could be represented by something like Factor and Logfactor funsors which can be used for pattern matching. Not sure if this adds anything new though.

As a side note, can Funsor be used to represent partially applied functions?

[fritzo]

@ordabayevy yes I think your Log functor relates the semiring pair ((add,mul), (logaddexp,add)) and the pair ((max,mul),(max,add)). However no functor relates e.g. the pair ((add,mul), (max,add)).

Indeed Funsors can represent functions, and funsor substitution represents partially applied functions.