torchtangents 0.1.1

TorchTangents (torchtangents) is a PyTorch-native port of Google's JAX-native Neural Tangents library. It lets you build, lift, and analyze deep neural networks in standard PyTorch while extracting their exact infinite-width Neural Network Gaussian Process (NNGP) and Neural Tangent Kernel (NTK) in closed form.

Novak et al., "Fast and Easy Infinite Neural Networks in Python" (ICLR 2020) introduced the original Neural Tangents library for JAX. With the original JAX repository archived in May 2025, torchtangents brings this mathematical framework natively to PyTorch — adding stateful torch.nn.Module wrappers, torch.fx model lifting, torch.func input Jacobians/Hessians, and torch.compile acceleration.

Installation

From PyPI

pip install torchtangents

Editable / Developer Installation

git clone https://github.com/deoliveirajoshua/torchtangents.git
cd torchtangents
pip install -e .

For running PyTorch-vs-JAX parity tests or benchmarks, see the Developer & Benchmark Installation Guide.

Requirements: Python >= 3.9, PyTorch >= 2.2.

30-Second Quickstart

import torch
import torch.nn as nn
import torchtangents as tt

x_train = torch.randn(16, 10)
y_train = torch.randn(16, 2).to(torch.float64)
x_test  = torch.randn(4, 10)

# 1. Build a dual-mode TorchTangents model (drop-in replacement for torch.nn)
model = tt.nn.Sequential(
    tt.nn.Linear(10, 128),
    tt.nn.ReLU(),
    tt.nn.Linear(128, 2),
)
out   = model(x_test)                              # Finite-width forward pass: (4, 2)
k_inf = model.inf_kernel(x_train, x_test)          # Infinite-width NNGP & NTK: (16, 4)
k_emp = model.empirical_ntk(x_train, x_test)       # Finite-width empirical NTK: (16, 4, 2, 2)

# 2. Or lift ANY existing vanilla PyTorch model in one line via torch.fx
vanilla_model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 2))
lifted_model  = tt.convert.lift(vanilla_model)
k_lifted      = lifted_model.inf_kernel(x_train, x_test, get="ntk")  # Shape: (16, 4)

# 3. Train an infinitely wide network in closed form (zero training loop!)
k_tr_tr = model.inf_kernel(x_train, None, get=("nngp", "ntk"))
k_te_tr = model.inf_kernel(x_test, x_train, get=("nngp", "ntk"))
gp_pred = tt.predict.gp_inference(k_tr_tr, y_train, diag_reg=1e-4)
y_pred  = gp_pred(get="ntk", k_test_train=k_te_tr).mean              # Shape: (4, 2)

Package Architecture

torchtangents is organized into four core submodules:

Theory at a Glance

1. Finite-Width (Empirical) NTK: Dot Product in Gradient Space

For a parameterized neural network $f(x; \theta) \in \mathbb{R}^K$ with finite width and parameters $\theta \in \mathbb{R}^P$, the Empirical Neural Tangent Kernel (model.empirical_ntk(x1, x2)) measures how a gradient step at input $x'$ changes the network's output at $x$. Mathematically, it is the inner product of the per-sample parameter Jacobians in gradient space:

$$\hat{\Theta}(x, x') = \left\langle \nabla_\theta f(x; \theta),, \nabla_\theta f(x'; \theta) \right\rangle = \sum_{p=1}^{P} \frac{\partial f(x; \theta)}{\partial \theta_p} \left(\frac{\partial f(x'; \theta)}{\partial \theta_p}\right)^\top \in \mathbb{R}^{K \times K}$$

At finite width, $\hat{\Theta}(x, x')$ depends on the random initialization of $\theta$ and evolves non-linearly during training.

2. Infinite-Width Limit: Layer-Wise Recursive NNGP & NTK

As the hidden layer widths approach infinity ($n_1, \dots, n_{L-1} \to \infty$) under NTK parameterization, two remarkable simplifications occur:

1. At Initialization (NNGP): Pre-activations at each layer $\ell$ converge by the Central Limit Theorem to an exact zero-mean Gaussian Process with covariance kernel $\Sigma^{(\ell)}(x, x')$.
2. During Training (NTK): The empirical kernel $\hat{\Theta}(x, x')$ converges deterministically to a scalar-output multiple $\Theta^{(L)}(x, x') \otimes I_K$ that remains constant (frozen) throughout gradient descent.

Both $\Sigma^{(L)}$ and $\Theta^{(L)}$ are computed analytically (model.inf_kernel(x1, x2)) via a closed-form layer-by-layer recursion:

• Base Case (First Linear Layer): $$\Sigma^{(1)}(x, x') = \frac{\sigma_w^2}{d_{\text{in}}} \langle x, x' \rangle + \sigma_b^2, \qquad \Theta^{(1)}(x, x') = \Sigma^{(1)}(x, x')$$

• Recursive Step (Layer $\ell \to \ell + 1$ with nonlinearity $\phi$): $$\Sigma^{(\ell+1)}(x, x') = \sigma_w^2 , \mathbb{E}_{(u, v) \sim \mathcal{N}\left(0,, \Lambda^{(\ell)}(x, x')\right)} \Big[\phi(u),\phi(v)\Big] + \sigma_b^2$$ $$\Theta^{(\ell+1)}(x, x') = \Sigma^{(\ell+1)}(x, x') + \sigma_w^2 , \dot{\Sigma}^{(\ell+1)}(x, x') , \Theta^{(\ell)}(x, x')$$

where the derivative covariance is

$$\dot{\Sigma}^{(\ell+1)}(x, x') = \mathbb{E}_{(u, v) \sim \mathcal{N}\left(0,, \Lambda^{(\ell)}(x, x')\right)} \Big[\dot{\phi}(u),\dot{\phi}(v)\Big], \qquad \Lambda^{(\ell)}(x, x') = \begin{pmatrix} \Sigma^{(\ell)}(x, x) & \Sigma^{(\ell)}(x, x') \ \Sigma^{(\ell)}(x', x) & \Sigma^{(\ell)}(x', x') \end{pmatrix}$$

3. Closed-Form Infinite-Width Training

Because $\Theta^{(L)}$ is constant during training in the infinite-width limit, training a deep network with continuous-time gradient descent under MSE loss reduces to linear ODE dynamics with an exact closed-form solution at $t \to \infty$:

$$\mu(x^{\star}) = \Theta^{(L)}(x^{\star}, X) \Big(\Theta^{(L)}(X, X) + \lambda I\Big)^{-1} Y$$

where $(X, Y)$ is the training dataset, $x^{\star}$ is the test input, and $\lambda$ is ridge regularization.

Core Usage Recipes

All examples below have standalone runnable scripts in examples/.

1. Stateful (tt.nn) vs. Stateless (tt.stax) Construction

See examples/1_model_construction.py:

import torch
import torchtangents as tt

x1, x2 = torch.randn(4, 10), torch.randn(2, 10)

# Option A: Stateful PyTorch-style module (finite forward + infinite kernel)
model = tt.nn.Sequential(tt.nn.Linear(10, 128), tt.nn.ReLU(), tt.nn.Linear(128, 5))
out  = model(x1)                   # Finite-width forward pass: (4, 5)
k_nn = model.inf_kernel(x1, x2)    # Infinite-width NNGP & NTK: (4, 2)

# Option B: Stateless functional stax (infinite kernel only, zero weights allocated)
kernel_fn = tt.stax.Sequential(tt.stax.Linear(128), tt.stax.ReLU(), tt.stax.Linear(5))
k_stax    = kernel_fn(x1, x2)      # Identical NNGP & NTK:      (4, 2)

2. Kernel Utilities: Empirical NTK, Input Gradients & Hessians

See examples/2_kernel_utilities.py:

x1 = torch.randn(2, 10, requires_grad=True)
x2 = torch.randn(3, 10)

# Analytical infinite-width kernel vs. finite-width empirical NTK
k_inf = model.inf_kernel(x1, x2)     # k_inf.nngp: (2, 3), k_inf.ntk: (2, 3)
k_emp = model.empirical_ntk(x1, x2)  # Shape: (2, 3, 5, 5) -> (N1, N2, Out, Out)

# First-order Jacobians and second-order Hessians of the infinite-width NTK w.r.t. x1
grad_pairs = model.ntk_gradient(x1, x2, mode="all_pairs")  # Shape: (2, 3, 10)
grad_gram  = model.ntk_gradient(x1, x2, mode="gram")       # Shape: (2, 10)
hess_pairs = model.ntk_hessian(x1, x2, mode="all_pairs")   # Shape: (2, 3, 10, 10)
hess_gram  = model.ntk_hessian(x1, x2, mode="gram")        # Shape: (2, 10, 2, 10)

3. Lifting Existing Vanilla PyTorch Models

See examples/3_model_lifting.py and examples/example_auto_converter.py:

import torch.nn as nn

class ResMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1  = nn.Linear(10, 16)
        self.fc2  = nn.Linear(16, 16)
        self.head = nn.Linear(16, 2)

    def forward(self, x):
        h = torch.relu(self.fc1(x))
        h = torch.relu(self.fc2(h)) + h  # Residual skip connection traces automatically
        return self.head(h)

# Lift into a dual-mode StaxModule
lifted = tt.convert.lift(ResMLP())
ntk    = lifted.inf_kernel(x1, x2, get="ntk")  # Shape: (2, 3)

# Or extract a lightweight stateless kernel function
kernel_fn = tt.convert.infinite_analog(ResMLP())
ntk       = kernel_fn(x1, x2, get="ntk")       # Shape: (2, 3)

4. Closed-Form Infinite-Width Training & GP Inference

See examples/5_infinite_width_prediction.py:

x_train = torch.randn(16, 8)
y_train = torch.sin(x_train[:, :1]).to(torch.float64)
x_test  = torch.randn(4, 8)

kernel_fn     = tt.stax.Sequential(tt.stax.Linear(256), tt.stax.ReLU(), tt.stax.Linear(1))
k_train_train = kernel_fn(x_train, None, get=("nngp", "ntk"))
k_test_train  = kernel_fn(x_test, x_train, get=("nngp", "ntk"))

# Exact Bayesian NNGP & NTK posterior mean
gp_predict = tt.predict.gp_inference(k_train_train, y_train, diag_reg=1e-4)
ntk_mean   = gp_predict(get="ntk", k_test_train=k_test_train).mean   # Shape: (4, 1)

# Exact continuous-time gradient descent trajectory at finite time t=1.0
gd_predict = tt.predict.gradient_descent_mse(k_train_train, y_train, diag_reg=1e-4)
fx_tr, fx_te = gd_predict(
    t=1.0,
    fx_train_0=torch.zeros_like(y_train),
    fx_test_0=torch.zeros(len(x_test), 1, dtype=torch.float64),
    k_test_train=k_test_train,
)

5. Mini-Batching & torch.compile Acceleration

See examples/6_batched_and_compiled_kernels.py:

# Compute exact kernels in mini-batches to fit large datasets in GPU memory
batched_kernel_fn = tt.stax.batch(kernel_fn, batch_size=8)
k_batched         = batched_kernel_fn(x_train, x_test, get=("nngp", "ntk"))

# Fuse and accelerate kernel evaluation with torch.compile (Linux / WSL2)
compiled_kernel_fn = tt.convert.optimize_model(kernel_fn)
k_compiled         = compiled_kernel_fn(x_train, x_test, get=("nngp", "ntk"))

Platform Support

• OS: Windows, Windows + WSL2, Linux
• Hardware: CPU, or CPU + NVIDIA GPUs (CUDA 12.x or later)
• Note: torch.compile (inductor backend) is supported on Linux and Windows + WSL2.

Performance

TorchTangents vs. Neural Tangents (JAX): time to calculate the infinite-width NTK of a WideResNet model.

import torchtangents.stax as tt_stax

def WideResnetBlock_tt(channels, strides=(1, 1), channel_mismatch=False):
    Main = tt_stax.serial(
        tt_stax.ReLU(), tt_stax.Conv(channels, (3, 3), strides, padding='SAME'),
        tt_stax.ReLU(), tt_stax.Conv(channels, (3, 3), padding='SAME'))
    Shortcut = tt_stax.Identity() if not channel_mismatch else tt_stax.Conv(
        channels, (3, 3), strides, padding='SAME')
    return tt_stax.serial(
        tt_stax.FanOut(2),
        tt_stax.parallel(Main, Shortcut),
        tt_stax.FanInSum())

def WideResnetGroup_tt(n, channels, strides=(1, 1)):
    blocks = [WideResnetBlock_tt(channels, strides, channel_mismatch=True)]
    for _ in range(n - 1):
        blocks.append(WideResnetBlock_tt(channels, (1, 1)))
    return tt_stax.serial(*blocks)

def WideResnet_tt(block_size, k, num_classes):
    return tt_stax.serial(
        tt_stax.Conv(16, (3, 3), padding='SAME'),
        WideResnetGroup_tt(block_size, int(16 * k)),
        WideResnetGroup_tt(block_size, int(32 * k), (2, 2)),
        WideResnetGroup_tt(block_size, int(64 * k), (2, 2)),
        tt_stax.AvgPool((8, 8)),
        tt_stax.Flatten(),
        tt_stax.Dense(num_classes, 1., 0.))

The speed % column is the proportion faster or slower torchtangents is compared to neural-tangents (higher is better). E.g., 1.0 = same speed, 1.2 = torchtangents is 20% faster.

Citation

If you use torchtangents in academic work, please cite:

@software{torchtangents2025,
  title        = {torchtangents: Infinite-width neural networks and NTKs in PyTorch},
  author       = {Joshua DeOliveira},
  year         = {2025},
  url          = {https://github.com/deoliveirajoshua/torchtangents}
}

Additionally, cite when appropriate these foundational works that torchtangents mirrors in PyTorch:

# Infinite width NTK/NNGP:
@inproceedings{neuraltangents2020,
    title={Neural Tangents: Fast and Easy Infinite Neural Networks in Python},
    author={Roman Novak and Lechao Xiao and Jiri Hron and Jaehoon Lee and Alexander A. Alemi and Jascha Sohl-Dickstein and Samuel S. Schoenholz},
    booktitle={International Conference on Learning Representations},
    year={2020},
    pdf={https://arxiv.org/abs/1912.02803},
    url={https://github.com/google/neural-tangents}
}

# Finite width, empirical NTK/NNGP:
@inproceedings{novak2022fast,
    title={Fast Finite Width Neural Tangent Kernel},
    author={Roman Novak and Jascha Sohl-Dickstein and Samuel S. Schoenholz},
    booktitle={International Conference on Machine Learning},
    year={2022},
    pdf={https://arxiv.org/abs/2206.08720},
    url={https://github.com/google/neural-tangents}
}

# Attention and variable-length inputs:
@inproceedings{hron2020infinite,
    title={Infinite attention: NNGP and NTK for deep attention networks},
    author={Jiri Hron and Yasaman Bahri and Jascha Sohl-Dickstein and Roman Novak},
    booktitle={International Conference on Machine Learning},
    year={2020},
    pdf={https://arxiv.org/abs/2006.10540},
    url={https://github.com/google/neural-tangents}
}

# Infinite-width "standard" parameterization:
@misc{sohl2020on,
    title={On the infinite width limit of neural networks with a standard parameterization},
    author={Jascha Sohl-Dickstein and Roman Novak and Samuel S. Schoenholz and Jaehoon Lee},
    publisher = {arXiv},
    year={2020},
    pdf={https://arxiv.org/abs/2001.07301},
    url={https://github.com/google/neural-tangents}
}

# Elementwise nonlinearities and sketching:
@inproceedings{han2022fast,
    title={Fast Neural Kernel Embeddings for General Activations},
    author={Insu Han and Amir Zandieh and Jaehoon Lee and Roman Novak and Lechao Xiao and Amin Karbasi},
    booktitle = {Advances in Neural Information Processing Systems},
    year={2022},
    pdf={https://arxiv.org/abs/2209.04121},
    url={https://github.com/google/neural-tangents}
}

License & Acknowledgments

TorchTangents operates under the open-source Apache License 2.0.

Because TorchTangents' stax sub-module (torchtangents.stax) mirrors and ports the core mathematical abstractions of Google's original Neural Tangents (JAX), we proudly acknowledge and extend the original computational logic developed by the Google team. For legal attributions, see the NOTICE and LICENSE files.