{"id":13687789,"url":"https://github.com/pytorch/functorch","last_synced_at":"2025-05-13T20:19:00.802Z","repository":{"id":37046658,"uuid":"359624909","full_name":"pytorch/functorch","owner":"pytorch","description":"functorch is JAX-like composable function transforms for PyTorch.","archived":false,"fork":false,"pushed_at":"2025-05-06T00:50:35.000Z","size":21382,"stargazers_count":1424,"open_issues_count":249,"forks_count":104,"subscribers_count":25,"default_branch":"main","last_synced_at":"2025-05-06T01:35:53.527Z","etag":null,"topics":["gradients","hessians","pytorch"],"latest_commit_sha":null,"homepage":"https://pytorch.org/functorch/","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"bsd-3-clause","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/pytorch.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2021-04-19T23:24:04.000Z","updated_at":"2025-04-25T02:04:22.000Z","dependencies_parsed_at":"2023-10-17T06:05:11.770Z","dependency_job_id":"eb562ed5-a6b6-4012-a638-ca7270b53a2b","html_url":"https://github.com/pytorch/functorch","commit_stats":{"total_commits":1315,"total_committers":62,"mean_commits":21.20967741935484,"dds":0.7072243346007605,"last_synced_commit":"b71aa0b4387b86c278132209b99538be48ef4c74"},"previous_names":["zou3519/functorch"],"tags_count":6,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pytorch%2Ffunctorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pytorch%2Ffunctorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pytorch%2Ffunctorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pytorch%2Ffunctorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/pytorch","download_url":"https://codeload.github.com/pytorch/functorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254020659,"owners_count":22000757,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["gradients","hessians","pytorch"],"created_at":"2024-08-02T15:01:00.622Z","updated_at":"2025-05-13T20:19:00.784Z","avatar_url":"https://github.com/pytorch.png","language":"Jupyter Notebook","funding_links":[],"categories":["Jupyter Notebook"],"sub_categories":[],"readme":"# functorch\n\n[**Why functorch?**](#why-composable-function-transforms)\n| [**Install guide**](#install)\n| [**Transformations**](#what-are-the-transforms)\n| [**Documentation**](#documentation)\n| [**Future Plans**](#future-plans)\n\n**This library is currently under heavy development - if you have suggestions\non the API or use-cases you'd like to be covered, please open an github issue\nor reach out. We'd love to hear about how you're using the library.**\n\n`functorch` is [JAX-like](https://github.com/google/jax) composable function\ntransforms for PyTorch.\n\nIt aims to provide composable `vmap` and `grad` transforms that work with\nPyTorch modules and PyTorch autograd with good eager-mode performance.\n\nIn addition, there is experimental functionality to trace through these\ntransformations using FX in order to capture the results of these transforms\nahead of time. This would allow us to compile the results of vmap or grad\nto improve performance.\n\n## Why composable function transforms?\n\nThere are a number of use cases that are tricky to do in\nPyTorch today:\n- computing per-sample-gradients (or other per-sample quantities)\n- running ensembles of models on a single machine\n- efficiently batching together tasks in the inner-loop of MAML\n- efficiently computing Jacobians and Hessians\n- efficiently computing batched Jacobians and Hessians\n\nComposing `vmap`, `grad`, `vjp`, and `jvp` transforms allows us to express the above\nwithout designing a separate subsystem for each. This idea of composable function\ntransforms comes from the [JAX framework](https://github.com/google/jax).\n\n## Install\n\nThere are two ways to install functorch:\n1. functorch from source\n2. functorch beta (compatible with recent PyTorch releases)\n\nWe recommend trying out the functorch beta first.\n\n### Installing functorch from source\n\n\u003cdetails\u003e\u003csummary\u003eClick to expand\u003c/summary\u003e\n\u003cp\u003e\n\n#### Using Colab\n\nFollow the instructions [in this Colab notebook](https://colab.research.google.com/drive/1CrLkqIrydBYP_svnF89UUO-aQEqNPE8x?usp=sharing)\n\n#### Locally\n\nAs of 9/21/2022, `functorch` comes installed alongside a nightly PyTorch binary.\nPlease install a Preview (nightly) PyTorch binary; see  https://pytorch.org/\nfor instructions.\n\nOnce you've done that, run a quick sanity check in Python:\n```py\nimport torch\nfrom functorch import vmap\nx = torch.randn(3)\ny = vmap(torch.sin)(x)\nassert torch.allclose(y, x.sin())\n```\n\n#### functorch development setup\n\nAs of 9/21/2022, `functorch` comes installed alongside PyTorch and is in the\nPyTorch source tree. Please install\n[PyTorch from source](https://github.com/pytorch/pytorch#from-source), then,\nyou will be able to `import functorch`.\n\nTry to run some tests to make sure all is OK:\n```bash\npytest test/test_vmap.py -v\npytest test/test_eager_transforms.py -v\n```\n\nAOTAutograd has some additional optional requirements. You can install them via:\n```bash\npip install networkx\n```\n\nTo run functorch tests, please install our test dependencies (`expecttest`, `pyyaml`).\n\n\n\u003c/p\u003e\n\u003c/details\u003e\n\n### Installing functorch beta (compatible with recent PyTorch releases)\n\n\u003cdetails\u003e\u003csummary\u003eClick to expand\u003c/summary\u003e\n\u003cp\u003e\n\n#### Using Colab\n\nFollow the instructions [here](https://colab.research.google.com/drive/1GNfb01W_xf8JRu78ZKoNnLqiwcrJrbYG#scrollTo=HJ1srOGeNCGA)\n\n#### pip\n\nPrerequisite: [Install PyTorch](https://pytorch.org/get-started/locally/)\n\n\n```bash\npip install functorch\n```\n\nFinally, run a quick sanity check in python:\n```py\nimport torch\nfrom functorch import vmap\nx = torch.randn(3)\ny = vmap(torch.sin)(x)\nassert torch.allclose(y, x.sin())\n```\n\n\u003c/p\u003e\n\u003c/details\u003e\n\n## What are the transforms?\n\nRight now, we support the following transforms:\n- `grad`, `vjp`, `jvp`,\n- `jacrev`, `jacfwd`, `hessian`\n- `vmap`\n\nFurthermore, we have some utilities for working with PyTorch modules.\n- `make_functional(model)`\n- `make_functional_with_buffers(model)`\n\n### vmap\n\nNote: `vmap` imposes restrictions on the code that it can be used on.\nFor more details, please read its docstring.\n\n`vmap(func)(*inputs)` is a transform that adds a dimension to all Tensor\noperations in `func`. `vmap(func)` returns a new function that maps `func` over\nsome dimension (default: 0) of each Tensor in `inputs`.\n\n`vmap` is useful for hiding batch dimensions: one can write a function `func`\nthat runs on examples and then lift it to a function that can take batches of\nexamples with `vmap(func)`, leading to a simpler modeling experience:\n\n```py\nfrom functorch import vmap\nbatch_size, feature_size = 3, 5\nweights = torch.randn(feature_size, requires_grad=True)\n\ndef model(feature_vec):\n    # Very simple linear model with activation\n    assert feature_vec.dim() == 1\n    return feature_vec.dot(weights).relu()\n\nexamples = torch.randn(batch_size, feature_size)\nresult = vmap(model)(examples)\n```\n\n### grad\n\n`grad(func)(*inputs)` assumes `func` returns a single-element Tensor. It compute\nthe gradients of the output of func w.r.t. to `inputs[0]`.\n\n```py\nfrom functorch import grad\nx = torch.randn([])\ncos_x = grad(lambda x: torch.sin(x))(x)\nassert torch.allclose(cos_x, x.cos())\n\n# Second-order gradients\nneg_sin_x = grad(grad(lambda x: torch.sin(x)))(x)\nassert torch.allclose(neg_sin_x, -x.sin())\n```\n\nWhen composed with `vmap`, `grad` can be used to compute per-sample-gradients:\n```py\nfrom functorch import vmap\nbatch_size, feature_size = 3, 5\n\ndef model(weights,feature_vec):\n    # Very simple linear model with activation\n    assert feature_vec.dim() == 1\n    return feature_vec.dot(weights).relu()\n\ndef compute_loss(weights, example, target):\n    y = model(weights, example)\n    return ((y - target) ** 2).mean()  # MSELoss\n\nweights = torch.randn(feature_size, requires_grad=True)\nexamples = torch.randn(batch_size, feature_size)\ntargets = torch.randn(batch_size)\ninputs = (weights,examples, targets)\ngrad_weight_per_example = vmap(grad(compute_loss), in_dims=(None, 0, 0))(*inputs)\n```\n\n### vjp\n\nThe `vjp` transform applies `func` to `inputs` and returns a new function that\ncomputes vjps given some `cotangents` Tensors.\n```py\nfrom functorch import vjp\noutputs, vjp_fn = vjp(func, inputs); vjps = vjp_fn(*cotangents)\n```\n\n### jvp\n\nThe `jvp` transforms computes Jacobian-vector-products and is also known as\n\"forward-mode AD\". It is not a higher-order function unlike most other transforms,\nbut it returns the outputs of `func(inputs)` as well as the `jvp`s.\n```py\nfrom functorch import jvp\nx = torch.randn(5)\ny = torch.randn(5)\nf = lambda x, y: (x * y)\n_, output = jvp(f, (x, y), (torch.ones(5), torch.ones(5)))\nassert torch.allclose(output, x + y)\n```\n\n### jacrev, jacfwd, and hessian\n\nThe `jacrev` transform returns a new function that takes in `x` and returns the\nJacobian of `torch.sin` with respect to `x` using reverse-mode AD.\n```py\nfrom functorch import jacrev\nx = torch.randn(5)\njacobian = jacrev(torch.sin)(x)\nexpected = torch.diag(torch.cos(x))\nassert torch.allclose(jacobian, expected)\n```\nUse `jacrev` to compute the jacobian. This can be composed with vmap to produce\nbatched jacobians:\n\n```py\nx = torch.randn(64, 5)\njacobian = vmap(jacrev(torch.sin))(x)\nassert jacobian.shape == (64, 5, 5)\n```\n\n`jacfwd` is a drop-in replacement for `jacrev` that computes Jacobians using\nforward-mode AD:\n```py\nfrom functorch import jacfwd\nx = torch.randn(5)\njacobian = jacfwd(torch.sin)(x)\nexpected = torch.diag(torch.cos(x))\nassert torch.allclose(jacobian, expected)\n```\n\nComposing `jacrev` with itself or `jacfwd` can produce hessians:\n```py\ndef f(x):\n  return x.sin().sum()\n\nx = torch.randn(5)\nhessian0 = jacrev(jacrev(f))(x)\nhessian1 = jacfwd(jacrev(f))(x)\n```\n\nThe `hessian` is a convenience function that combines `jacfwd` and `jacrev`:\n```py\nfrom functorch import hessian\n\ndef f(x):\n  return x.sin().sum()\n\nx = torch.randn(5)\nhess = hessian(f)(x)\n```\n\n### Tracing through the transformations\nWe can also trace through these transformations in order to capture the results as new code using `make_fx`. There is also experimental integration with the NNC compiler (only works on CPU for now!).\n\n```py\nfrom functorch import make_fx, grad\ndef f(x):\n    return torch.sin(x).sum()\nx = torch.randn(100)\ngrad_f = make_fx(grad(f))(x)\nprint(grad_f.code)\n\ndef forward(self, x_1):\n    sin = torch.ops.aten.sin(x_1)\n    sum_1 = torch.ops.aten.sum(sin, None);  sin = None\n    cos = torch.ops.aten.cos(x_1);  x_1 = None\n    _tensor_constant0 = self._tensor_constant0\n    mul = torch.ops.aten.mul(_tensor_constant0, cos);  _tensor_constant0 = cos = None\n    return mul\n```\n\n### Working with NN modules: make_functional and friends\n\nSometimes you may want to perform a transform with respect to the parameters\nand/or buffers of an nn.Module. This can happen for example in:\n- model ensembling, where all of your weights and buffers have an additional\ndimension\n- per-sample-gradient computation where you want to compute per-sample-grads\nof the loss with respect to the model parameters\n\nOur solution to this right now is an API that, given an nn.Module, creates a\nstateless version of it that can be called like a function.\n\n- `make_functional(model)` returns a functional version of `model` and the\n`model.parameters()`\n- `make_functional_with_buffers(model)` returns a functional version of\n`model` and the `model.parameters()` and `model.buffers()`.\n\nHere's an example where we compute per-sample-gradients using an nn.Linear\nlayer:\n\n```py\nimport torch\nfrom functorch import make_functional, vmap, grad\n\nmodel = torch.nn.Linear(3, 3)\ndata = torch.randn(64, 3)\ntargets = torch.randn(64, 3)\n\nfunc_model, params = make_functional(model)\n\ndef compute_loss(params, data, targets):\n    preds = func_model(params, data)\n    return torch.mean((preds - targets) ** 2)\n\nper_sample_grads = vmap(grad(compute_loss), (None, 0, 0))(params, data, targets)\n```\n\nIf you're making an ensemble of models, you may find\n`combine_state_for_ensemble` useful.\n\n## Documentation\n\nFor more documentation, see [our docs website](https://pytorch.org/functorch).\n\n## Debugging\n`torch._C._functorch.dump_tensor`: Dumps dispatch keys on stack\n`torch._C._functorch._set_vmap_fallback_warning_enabled(False)` if the vmap warning spam bothers you.\n\n## Future Plans\n\nIn the end state, we'd like to upstream this into PyTorch once we iron out the\ndesign details. To figure out the details, we need your help -- please send us\nyour use cases by starting a conversation in the issue tracker or trying our\nproject out.\n\n## License\nFunctorch has a BSD-style license, as found in the [LICENSE](LICENSE) file.\n\n## Citing functorch\n\nIf you use functorch in your publication, please cite it by using the following BibTeX entry.\n\n```bibtex\n@Misc{functorch2021,\n  author =       {Horace He, Richard Zou},\n  title =        {functorch: JAX-like composable function transforms for PyTorch},\n  howpublished = {\\url{https://github.com/pytorch/functorch}},\n  year =         {2021}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpytorch%2Ffunctorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpytorch%2Ffunctorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpytorch%2Ffunctorch/lists"}