Metadata-Version: 2.4
Name: jet-for-pytorch
Version: 0.0.1
Summary: Taylor-mode automatic differentiation for PyTorch
Author: Felix Dangel, Tim Siebert, Marius Zeinhofer, Andrea Walther
License: MIT
Project-URL: Repository, https://github.com/f-dangel/torch-jet
Classifier: Development Status :: 3 - Alpha
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.12
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: einops
Requires-Dist: numpy
Requires-Dist: scipy
Requires-Dist: pyyaml
Provides-Extra: exp
Requires-Dist: memory_profiler; extra == "exp"
Requires-Dist: pandas; extra == "exp"
Requires-Dist: matplotlib; extra == "exp"
Requires-Dist: tueplots; extra == "exp"
Requires-Dist: jax; extra == "exp"
Requires-Dist: jaxlib; extra == "exp"
Requires-Dist: folx; extra == "exp"
Requires-Dist: arxiv-collector; extra == "exp"
Provides-Extra: test
Requires-Dist: coveralls; extra == "test"
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Provides-Extra: lint
Requires-Dist: ruff; extra == "lint"
Provides-Extra: doc
Requires-Dist: matplotlib; extra == "doc"
Requires-Dist: mkdocs==1.6.1; extra == "doc"
Requires-Dist: mkdocs-material==9.6.13; extra == "doc"
Requires-Dist: mkdocstrings[python]==0.30.0; extra == "doc"
Requires-Dist: mkdocs-gallery==0.10.4; extra == "doc"
Requires-Dist: pydot; extra == "doc"
Requires-Dist: pytest; extra == "doc"
Requires-Dist: tueplots; extra == "doc"
Requires-Dist: ruff; extra == "doc"
Dynamic: license-file

# Taylor Mode Autodiff in PyTorch

This library provides a PyTorch implementation of Taylor mode automatic differentiation, a generalization of forward mode to higher-order derivatives.
It is similar to JAX's Taylor mode ([`jax.experimental.jet`](https://docs.jax.dev/en/latest/jax.experimental.jet.html)).

The repository also hosts the Python functionality+experiments and LaTeX source for our NeurIPS 2025 paper ["Collapsing Taylor Mode Automatic Differentiation"](https://openreview.net/forum?id=XgQVL1uP34), which allows to further accelerate Taylor mode for many practical differential operators.

> 🔪 **Warning: expect rough edges!** 🔪
>
> This is a research prototype with various limitations (e.g. operator coverage).
> We highly recommend double-checking your results with PyTorch's autodiff.
> Please help us improve the package by providing feedback, filing issues, and opening pull requests.


## Getting Started

### Installation

```bash
pip install jet-for-pytorch
```

### Examples

See the [documentation](https://torch-jet.readthedocs.io/en/latest/generated/gallery/).

## Citing

If you find the `jet` package useful for your research, consider citing

```bibtex

@inproceedings{dangel2025collapsing,
  title =        {Collapsing Taylor Mode Automatic Differentiation},
  author =       {Felix Dangel and Tim Siebert and Marius Zeinhofer and Andrea
                  Walther},
  year =         2025,
  booktitle =    {Advances in Neural Information Processing Systems (NeurIPS)},
}

```
