원문 캡처 · github.com
원문 캡처 · github.com

x^4의 세 번째 도함수는 24x이며, 예제는 x=2에서 이를 계산한다. jet-for-pytorch는 테일러 모드 자동 미분으로 이 계산을 수행한다.

이 라이브러리는 파이토치(PyTorch)에서 고계 도함수를 구할 수 있게 한다. JAX의 jax.experimental.jet과 비슷하며, 미분 연산자를 더 빠르게 처리하는 방법을 다룬 NeurIPS 2025 논문의 기능과 실험 코드, LaTeX 소스도 저장소에 포함돼 있다.

테일러 모드 자동 미분은 무엇인가?

테일러 모드 자동 미분은 순방향 자동 미분(forward mode)을 고계 도함수까지 확장한 방법이다. 여러 차수의 입력 값을 함수에 전달해 이에 대응하는 미분 결과를 계산한다. 여러 차례의 도함수 계산을 한꺼번에 묶는 방식에 빗대면 이해하기 쉽다.

어떻게 설치하고 쓰나?

파이썬 패키지를 설치한다.

bash
pip install jet-for-pytorch

예제는 x^4 함수의 세 번째 도함수를 계산한다. 입력값 x₀에는 x를 넣고, x₁에는 1, x₂와 x₃에는 0을 넣으면 출력의 세 번째 항 f₃가 세 번째 도함수가 된다.

python
from torch import tensor, ones_like, zeros_like
from jet import jet

f = lambda x: x**4 # f'''(x) = 24 * x

x = tensor(2.0)
f_jet = jet(f, (x,)) # (x₀, x₁, x₂, x₃) ↦ (f₀, f₁, f₂, f₃)

# x₀ = x, x₁ = 1, x₂ = x₃ = 0 ⇒ f₃ = f'''(x)
f0, f1, f2, f3 = f_jet((x, ones_like(x), zeros_like(x), zeros_like(x)))

assert f3.allclose(24 * x) # third derivative matches by hand

지원 연산자 전체 목록은 입문 튜토리얼의 지원 연산자 절에서 확인할 수 있으며, 추가 예제는 문서에 있다. 지원하는 연산자 범위는 아직 확대 중이다. 저장소는 사용자의 피드백과 이슈 제보, 풀 리퀘스트를 요청한다.

논문을 인용할 때 사용할 정보는 다음과 같다.

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)},
}

원문은 지원 연산자의 전체 목록이나 설치에 필요한 파이토치 버전을 밝히지 않는다.