torchdiffeq

library0.2.5pypypi✓ verified 26d ago

torchdiffeq is a Python library providing ordinary differential equation (ODE) solvers implemented in PyTorch. It supports backpropagation through ODE solutions using the adjoint method, ensuring constant memory cost. The library offers a clean API for usage in deep learning applications, fully supporting GPU execution. The current version is 0.2.5, last released in November 2024, indicating an active development and maintenance cadence.

ai-ml

pip install torchdiffeq

github ↗package ↗

INSTALL

IMPORT

SIG · TORCHDIFFEQ

torchdiffeq

ai-mlpythonv0.2.5

Install

69.9s avg

Import

8160ms

Disk

4941MB

Pass rate

4/ 10

Env Coverage4 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.2.5 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

✕ build_error

✓ 79.3s

py 3.11

✕ build_error

✓ 73.8s

py 3.12

✕ build_error

✓ 65s

py 3.13

✕ build_error

✓ 61.4s

py 3.9

✕ build_error

✕ timeout

4941MB installed

● package 4941MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

odeint

✓ from torchdiffeq import odeint

Standard ODE solver for direct backpropagation.

odeint_adjoint

✓ from torchdiffeq import odeint_adjoint as odeint

✗ from torchdiffeq import odeint_adjoint

The common pattern is to alias `odeint_adjoint` to `odeint` for easy switching. When using `odeint_adjoint`, the ODE function (`func`) *must* be an `nn.Module` to collect parameters.

This quickstart demonstrates how to define a simple ODE function as an `nn.Module` and use `torchdiffeq.odeint` to solve it over a specified time interval. The output `solution` tensor contains the evaluated states at each time point.

import torch import torch.nn as nn from torchdiffeq import odeint # Define the ODE function as an nn.Module class ODEFunc(nn.Module): def forward(self, t, y): # Example ODE: dy/dt = -0.1y + t # y and t are torch.Tensor return -0.1 * y + t # Initial state y(t=0) y0 = torch.tensor([0.7]) # Time points at which to evaluate the solution t = torch.linspace(0., 10., 100) # 100 points from t=0 to t=10 # Solve the ODE using the default (dopri5) solver solution = odeint(ODEFunc(), y0, t) print("Shape of solution (time_steps, initial_dim):") print(solution.shape) # Expected: (100, 1) print("\nFirst 5 values of the solution:") print(solution[:5])

Debug

Known issues

gotchaWhen using `odeint_adjoint` for O(1) memory backpropagation, the ODE function (`func`) must be an instance of `torch.nn.Module`. This is crucial for the adjoint method to correctly identify and collect parameters for gradient computation.

fix

Ensure your ODE dynamics `func` inherits from `torch.nn.Module` if you intend to use `odeint_adjoint`.

affects: All versions

gotchaDirect backpropagation through `odeint` (without `odeint_adjoint`) can be memory-intensive, especially for complex ODE trajectories or long integration times, as it stores all intermediate states. For O(1) memory cost, use the adjoint method (`odeint_adjoint`).

fix

For memory-efficient training, import and use `odeint_adjoint` (often aliased as `odeint`) instead of the default `odeint`.

affects: All versions

gotchaAdaptive ODE solvers (like the default `dopri5`) use `rtol` (relative tolerance) and `atol` (absolute tolerance) to control the accuracy and number of steps. Incorrectly set tolerances can lead to either excessively slow computations or inaccurate solutions.

fix

Experiment with `rtol` and `atol` (e.g., `odeint(..., rtol=1e-3, atol=1e-5)`) to find a balance between speed and desired accuracy for your specific problem. Higher values mean faster but less accurate solutions.

affects: All versions

gotchaThe `dtype` for timelike quantities in solvers defaults to `torch.float64`. While more stable, using `torch.float32` can significantly improve speed but might lead to numerical instability or underflow issues in certain scenarios.

fix

Consider setting `options={'dtype': torch.float32}` within the `odeint` call if you need higher performance and have verified numerical stability with single-precision floats.

affects: All versions

Errors

Common errors & fixes

NotImplementedError: UserDefinedObjectVariable(EnumMeta) is not a constant

This error occurs when using torch.compile() on models that incorporate torchdiffeq for ODE solving, due to TorchDynamo's inability to handle enum types from torchdiffeq during compilation.

fix

As a workaround, avoid using torch.compile() with torchdiffeq until support for enum types is implemented in TorchDynamo.

ImportError: numpy.core.multiarray failed to import

This error can occur when there's a version mismatch between PyTorch and NumPy, leading to compatibility issues.

fix

Ensure that both PyTorch and NumPy are updated to compatible versions. For example, updating NumPy to version 1.19.3 has resolved similar issues in the past.

AttributeError: module 'torchdiffeq' has no attribute 'odeint_adjoint'

This error occurs when attempting to use the adjoint method for backpropagation, but the 'odeint_adjoint' function is not found in the torchdiffeq module.

fix

Ensure that you have the latest version of torchdiffeq installed, as 'odeint_adjoint' is available in version 0.2.5 and later.

TypeError: 'NoneType' object is not callable

This error can occur if the ODE function passed to 'odeint' is not properly defined or is set to None.

fix

Verify that the ODE function is correctly implemented and is not None before passing it to 'odeint'.

RuntimeError: CUDA error: device-side assert triggered

This error indicates that a CUDA assertion failed on the device side, often due to invalid operations or out-of-bounds memory access.

fix

Check the ODE function and input data for potential issues, and consider running the code on the CPU to get more informative error messages.

Upgrade

Version history

0.2.5latest on PyPI · released Nov 21, 2024

Audit

Dependencies

torchrequiredCore deep learning framework dependency.

scipyrequiredUsed for additional solver wrappers and numerical utilities.

Agent activity

36 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/rtqichen/torchdiffeq ↗

packagehttps://pypi.org/project/torchdiffeq/ ↗