• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1functorch
2=========
3
4.. currentmodule:: functorch
5
6.. warning::
7
8   We've integrated functorch into PyTorch. As the final step of the
9   integration, the functorch APIs are deprecated as of PyTorch 2.0.
10   Please use the torch.func APIs instead and see the
11   `migration guide <https://pytorch.org/docs/main/func.migrating.html>`_
12   and `docs <https://pytorch.org/docs/main/func.html>`_
13   for more details.
14
15Function Transforms
16-------------------
17.. autosummary::
18    :toctree: generated
19    :nosignatures:
20
21    vmap
22    grad
23    grad_and_value
24    vjp
25    jvp
26    jacrev
27    jacfwd
28    hessian
29    functionalize
30
31Utilities for working with torch.nn.Modules
32-------------------------------------------
33
34In general, you can transform over a function that calls a ``torch.nn.Module``.
35For example, the following is an example of computing a jacobian of a function
36that takes three values and returns three values:
37
38.. code-block:: python
39
40    model = torch.nn.Linear(3, 3)
41
42    def f(x):
43        return model(x)
44
45    x = torch.randn(3)
46    jacobian = jacrev(f)(x)
47    assert jacobian.shape == (3, 3)
48
49However, if you want to do something like compute a jacobian over the parameters
50of the model, then there needs to be a way to construct a function where the
51parameters are the inputs to the function.
52That's what :func:`make_functional` and :func:`make_functional_with_buffers` are for:
53given a ``torch.nn.Module``, these return a new function that accepts ``parameters``
54and the inputs to the Module's forward pass.
55
56.. autosummary::
57    :toctree: generated
58    :nosignatures:
59
60    make_functional
61    make_functional_with_buffers
62    combine_state_for_ensemble
63
64If you're looking for information on fixing Batch Norm modules, please follow the
65guidance here
66
67.. toctree::
68   :maxdepth: 1
69
70   batch_norm
71