Autodiff implementation (experimental)

[ordabayevy]

This is an implementation of autodiff. The goal is to address issues in computing expectations in TraceEnum_ELBO and TraceMarkovEnum_ELBO (#493). As of now it seems to fix nan gradients under eager interpretation in TraceEnum_ELBO.

The algorithm implements equivalents of linearize(), transpose() functions, and is tape-free (#446).

1. Linearize. Variables that need to be linearized are replaced by primal- tangent tuple JVP(primal, tangent) and then pattern matched to propagate tangents, e.g.:

JVP(x, dx) + JVP(y, dy) = JVP(x+y, dx+dy)
JVP(x, dx) * JVP(y, dy) = JVP(x*y, ydx + xdy)
JVP(x, dx) * y = JVP(x*y, ydx)

Out tangent is a linear function of in tangents. JVP is used for (add,mul) semiring and LJVP is used for (logaddexp,add) semiring.

2. Transpose of a linear function. Transpose is implemented simply by inverting the order of function execution and transposing matrices, in this case swapping more primitive operations .reduce(sum_op, "i") and .expand("i") (broadcasting does this automatically).

[fritzo]

should this be

fresh = frozenset(v.name for v in expanded_vars)

[fritzo]

test files should read but not write the global backend. Instead you can decorate each test with

@pytest.mark.skipif(get_backend() != "torch", reason="backend-specific")

and then run tests with

FUNSOR_BACKEND=torch pytest test/test_autodiff.py

[fritzo]

Hmm I'd like to better understand the need for this.

We've been trying to preserve the extensionality property in Funsor, which states that: if under every grounding substitution subs a pair of funsors f,g satisfy f(**subs) == g(**subs), then it should be permissible for an optimizer to replace funsor f with funsor g in any expression. IIUC this Expand funsor would break extensionality because f.expand(...) behaves as f under every grounding substitution.