Automatic Differentiation

[JonathanWenger]

How can we support or include automatic differentiation into ProbNum ?

There are potentially 3 scenarios:

1. Include a specific automatic differentiation framework (e.g PyTorch) via a dependency
2. Write a backward pass for every probabilistic numerical method
3. Write an interface (kernel?) for arbitrary automatic differentiation frameworks.

ad 1)
Write code using autodiff structures (tensors).

Advantages:

• can automatically provide speedup by producing compiled code
• gradients are computed automatically
  Disadvantages:
• have to support changes in the dependency into the future
• bound to one autodiff framework, change almost certainly necessitates package rewrite

ad 2)
Implement the backward pass for every probabilistic numerical method.

Advantages:

• able to support any autodiff framework
  Disadvantages
• gradients may be very difficult to compute analytically
• no free speedup via compiled code or GPU

ad 3)
Write an interface (kernel?) for arbitrary automatic differentiation frameworks.

Advantages:

• able to support any autodiff framework in principle
• no need to compute gradients by hand
  Disadvantages:
• no known example of major packages supporting different automatic differentiation frameworks
• interface has to be maintained for different frameworks
• no free speedup via compiled code or GPU

TODO:

• code up a mini working example of a blackbox function which provides a .backward() method for pytorch and tensorflow
• how do frameworks like JAX perform for linear algebra

Other

• Framework interfacing autograd, pytorch and tensorflow: https://github.com/geomstats/geomstats
• Tensor library with multiple backends (Jax, numpy, pytorch, tensorflow): https://github.com/tensorly/tensorly

[JonathanWenger]

Tensor library with multiple backends from developers close to Tübingen: eagerpy