Install & Compatibility
Where this runs
tested against v0.2.6 · 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
py 3.10
✕ build_error
✓ 83.2s
py 3.11
✕ build_error
✓ 74.1s
py 3.12
✕ build_error
✓ 64.7s
py 3.13
✕ build_error
✓ 59.7s
py 3.9
✕ build_error
✕ timeout
4941MB installed
● package 4941MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
sdeint
✓ from torchsde import sdeint
sdeint_adjoint
✓ from torchsde import sdeint_adjoint
Use this for memory-efficient backpropagation through SDEs, especially for long trajectories.
BrownianPath
✓ from torchsde.brownian_lib import BrownianPath
✗ from torchsde.brownian import BrownianPath
For optimal performance, use the C++ backend Brownian motion classes from `brownian_lib` (since v0.1.1).
This quickstart demonstrates how to define a simple SDE model by implementing the drift `f` and diffusion `g` functions as part of a `torch.nn.Module`. It then uses `torchsde.sdeint` to solve the SDE over a given time interval `ts` starting from an initial state `y0`. For optimal performance, defining `f_and_g` to compute both drift and diffusion simultaneously is recommended.
import torch
import torchsde
class SDE(torch.nn.Module):
def __init__(self, d, m):
super().__init__()
self.mu = torch.nn.Linear(d, d)
self.sigma = torch.nn.Linear(d, m)
def f(self, t, y):
return self.mu(y)
def g(self, t, y):
return self.sigma(y)
def f_and_g(self, t, y):
return self.mu(y), self.sigma(y)
# Define parameters
D = 2 # State dimension
M = 2 # Noise dimension
T = 1.0 # End time
sde = SDE(D, M)
ts = torch.linspace(0, T, 10)
eps = 0.1 # Small initial perturbation
y0 = torch.rand(1, D) * eps # Initial state
# Solve the SDE
with torch.no_grad(): # For inference, use no_grad
ys = torchsde.sdeint(sde, y0, ts)
print("SDE solved, output shape:", ys.shape)
# Expected output shape: (len(ts), batch_size, D)
# e.g., (10, 1, 2)
Debug
Known issues
breakingOlder versions (pre-0.2.6) of torchsde had known dependency resolution issues on PyPI, leading to installation failures or incorrect dependency versions.fixUpgrade to torchsde v0.2.6 or later to benefit from corrected dependency lists and improved CI. `pip install --upgrade torchsde`
affects: <0.2.6
gotchaThe choice between `sdeint` and `sdeint_adjoint` is crucial for memory efficiency during training. `sdeint_adjoint` uses adjoint sensitivity analysis to compute gradients with O(1) memory cost with respect to the SDE trajectory length, which is vital for long simulations.fixFor training deep learning models involving SDEs, always prefer `sdeint_adjoint` when backpropagating through the SDE, unless you have explicit reasons or memory is not a concern.
affects: All versions
gotchaFor optimal performance with Brownian motion generation, ensure you are importing `BrownianPath` or `BrownianTree` from `torchsde.brownian_lib` (the C++ backend) rather than `torchsde.brownian` (the older Python backend).fixChange your import statement from `from torchsde.brownian import BrownianPath` to `from torchsde.brownian_lib import BrownianPath`.
affects: All versions since v0.1.1
gotchatorchsde is tightly integrated with PyTorch and specific PyTorch versions. Incompatibility between torchsde and your installed PyTorch version can lead to runtime errors or unexpected behavior.fixAlways check the `requires_python` and `install_requires` in `setup.py` or PyPI for `torchsde` to ensure your PyTorch version meets the minimum requirements. Currently, `torchsde==0.2.6` requires `torch>=1.6.0`.
affects: All versions
Errors
Common errors & fixes
ERROR: Failed building wheel for torchsde
The `torchsde` package requires compiling custom C++/CUDA extensions during installation, which often fails if the necessary build tools (like `gcc`, `nvcc` for CUDA) are missing, or if there's an incompatibility with the installed PyTorch or CUDA version.
fixEnsure `gcc` (and `nvcc` if using GPU) is installed and in your system's `PATH`. Verify your PyTorch and CUDA versions are compatible with `torchsde`. You might also try `pip install torchsde --no-build-isolation` as a workaround for some build environment issues.
ModuleNotFoundError: No module named 'torchsde'
The `torchsde` package is not installed in the current Python environment or the environment where it was installed is not active.
fixInstall the package using pip: `pip install torchsde`.
RuntimeError: Only tensors that are dtypes of floating point or complex are supported
The `torchsde` solvers and internal operations expect input tensors (e.g., the initial state `y0` or tensors within the drift `f` and diffusion `g` functions) to be of a floating-point or complex data type, not integers or booleans.
fixEnsure all relevant tensors are cast to a supported floating-point dtype, for example: `y0 = y0.float()` or `y0 = y0.double()`.
RuntimeError: The size of tensor a (X) must match the size of tensor b (Y) at non-singleton dimension Z
This error often occurs when the output shapes of your drift function `f(t, y)` or diffusion function `g(t, y)` do not match the expected dimensions. The `f` function must return a tensor with the same shape as `y`, and the `g` function must return a tensor of shape `(*y.shape, noise_dim)`.
fixCarefully check the `f` and `g` functions to ensure they return tensors with the correct shapes. For example, if `y` has shape `(batch_size, state_dim)`, `f` should return `(batch_size, state_dim)`, and `g` should return `(batch_size, state_dim, noise_dim)`.
Upgrade
Version history
0.2.6latest on PyPI · released Sep 26, 2023
Audit
Dependencies
torchrequiredCore deep learning framework dependency; torchsde builds on PyTorch tensors and autograd.