functorch
Inspired by Google JAX, functorch is a library that offers composable vmap (vectorization) and autodiff transforms. It enables advanced autodiff use cases that would otherwise be tricky to express in PyTorch. Examples of these include:
- running ensembles of models on a single machine
- efficiently computing Jacobians and Hessians
- computing per-sample-gradients (or other per-sample quantities)
We’re excited to announce functorch 0.2.0 with a number of improvements and new experimental features.
Significantly improved coverage
We significantly improved coverage for functorch.jvp (our forward-mode autodiff API) and other APIs that rely on it (functorch.{jacfwd, hessian}).
(Prototype) functorch.experimental.functionalize