torchode

Compared to other PyTorch ODE solvers, torchode provides two primary advantages:

1. JIT Compatibility: torchode is fully compatible with PyTorch JIT compilation. This allows you to compile your entire model (including the ODE solver) for faster execution. Compiled models can also participate in advanced PyTorch features like ONNX export and optimization with nvFuser.

2. Batch Parallelization: Unlike torchdiffeq and TorchDyn, which treat a batch of $n$ samples as a single large ODE system, torchode tracks the solver state for each sample in the batch independently. This means:

   • Each sample can accept or reject steps independently.
   • The error, number of steps, and gradients for sample x are unaffected by whether it was batched with sample y or z.
   • This is particularly important when working with stiff equations or when you need to eliminate unintended interactions between samples in a batch.

torchode is optimized for 2-dimensional data with the shape batch x features. It does not natively support arbitrary data dimensions (e.g., images or graphs).

To use non-vector data, you must manually manage the shapes by:

1. Flattening your input data into a 2D batch x features tensor before passing it to torchode.
2. Reshaping the data back to its original dimensions (e.g., B x C x W x H for images) inside your ODE function f.
3. Flattening the output of your ODE function back to 2D before returning it to the solver.

# Example with a batch of RGB images
b, c, w, h = 8, 3, 72, 72
y0 = torch.randn(b, w, h)

def f(t, y, args):
    c, w, h = args

    # Reshape 2D to B x C x W x H tensor and apply convolutional neural network
    y = some_cnn(y.reshape((-1, c, w, h)))

    # Flatten back to 2D when we return control back to torchode
    return y.flatten(start_dim=1)

term = to.ODETerm(f, with_args=True)
problem = to.InitialValueProblem(y0=y0.flatten(start_dim=1), t_eval=..)
sol = solver.solve(problem, args=(c, w, h))

To solve an Initial Value Problem (IVP), you need to compose several components: an ODETerm defining the dynamics, a step_method (like Dopri5), a step_size_controller (like IntegralController), and a solver wrapper (like AutoDiffAdjoint). For optimal performance, it is recommended to use torch.compile on the solver. The solver handles batches independently, allowing different samples in a batch to take different step sizes based on their individual stiffness.

import matplotlib.pyplot as pp
import torch
import torchode as to

def f(t, y):
    return -0.5 * y

y0 = torch.tensor([[1.2], [5.0]])
n_steps = 10
t_eval = torch.stack((torch.linspace(0, 5, n_steps), torch.linspace(3, 4, n_steps)))

term = to.ODETerm(f)
step_method = to.Dopri5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
solver = to.AutoDiffAdjoint(step_method, step_size_controller)
jit_solver = torch.compile(solver)

sol = jit_solver.solve(to.InitialValueProblem(y0=y0, t_eval=t_eval))
print(sol.stats)
# => {'n_f_evals': tensor([26, 26]), 'n_steps': tensor([4, 2]),
# =>  'n_accepted': tensor([4, 2]), 'n_initialized': tensor([10, 10])}

pp.plot(sol.ts[0], sol.ys[0])
pp.plot(sol.ts[1], sol.ys[1])

When computing gradients of ODE dynamics (e.g., in Neural ODEs) with respect to the solution, torchode provides three adjoint methods. Your choice depends on your memory constraints and model compatibility:

1. to.AutoDiffAdjoint (Preferred): Uses PyTorch's autograd to backpropagate directly through the solver's computation graph. It is fast and accurate but requires significant memory as it stores the entire forward pass.
2. to.BacksolveAdjoint (Memory Efficient): Solves the adjoint equations backwards in time. It uses torch.func to compute model gradients, making it much more memory-efficient than AutoDiffAdjoint.
3. to.JointBacksolveAdjoint (Fallback/Last Resort): Uses standard PyTorch autograd for model gradients. Use this only if your model is incompatible with torch.func.

Important caveats for to.JointBacksolveAdjoint:

• It solves the batch of adjoint equations as a single joint system, which introduces a subtle coupling between ODE solutions in a batch, breaking torchode's principle of complete independence.
• It is only applicable if all ODEs in your batch share the same evaluation points.

To solve problems where each instance in a batch requires unique parameters (e.g., individual decay rates), you can pass arbitrary objects to your dynamics function f.

1. Instantiate to.ODETerm with the argument with_args=True. This changes the expected signature of your dynamics function from f(t, y) to f(t, y, extra_args).
2. When calling solver.solve(), pass your parameters via the args keyword argument.

Important: torchode does not inspect or manage the contents of args. You are responsible for ensuring that the extra_args passed to your function are correctly shaped for broadcasting against your state y and time t within the dynamics function.

import torch
import torchode as to

# The dynamics function now accepts a third argument
def f(t, y, decay_rate):
    # You must handle broadcasting (e.g., using [:, None]) manually
    return -decay_rate[:, None] * y

# Enable extra arguments in ODETerm
term = to.ODETerm(f, with_args=True)
step_method = to.Tsit5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
solver = to.AutoDiffAdjoint(step_method, step_size_controller)

# Setup problem
y0 = torch.tensor([[1.2], [5.0]])
t_eval = torch.linspace(0, 5, 10).unsqueeze(1).repeat(1, 2).T # Example t_eval
decay_rate = torch.tensor([0.1, 5.0]) # Parameters for each batch instance

problem = to.InitialValueProblem(y0=y0, t_eval=t_eval)

# Pass the parameters via the 'args' parameter
sol = solver.solve(problem, args=decay_rate)

When choosing an adjoint method in torchode, consider the following mental models:

1. AutoDiffAdjoint: Uses standard PyTorch autograd. It is the simplest to use and works with any nn.Module that defines a vector field. However, it stores all intermediate states of the forward pass to compute gradients, which can be very memory-intensive for long integrations.

2. BacksolveAdjoint: Implements the adjoint sensitivity method. It solves an augmented system backwards in time. This is significantly more memory-efficient because it doesn't need to store the entire forward trajectory. It uses torch.func for high-performance VJP computations.

3. JointBacksolveAdjoint: An optimization of the BacksolveAdjoint for batches where instances have different time domains. It uses a linear mapping (unwrapping) to treat the batch as a single integration problem, improving computational efficiency through better batching.