### Setup Dummy DataLoader and Time Span

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2d_generalized_adjoint.ipynb

Creates a dummy training dataloader and defines the time span for integration. This setup is required for compatibility with PyTorch Lightning's API, even when sampling initial conditions dynamically.

```python
# dummy trainloader
train = torch.utils.data.TensorDataset(torch.zeros(1), torch.zeros(1))
trainloader = torch.utils.data.DataLoader(train, batch_size=1, shuffle=True)
t_span = torch.linspace(0, 2, 500)
```

--------------------------------

### Import Libraries and Setup

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2d_generalized_adjoint.ipynb

Imports necessary modules from torchdyn and PyTorch, and sets up autoreload for notebooks. It also includes a flag for quick validation runs.

```python
from torchdyn.core import NeuralODE
from torchdyn.nn import *
from torchdyn.datasets import *
from torchdyn.utils import *

%load_ext autoreload
%autoreload 2

# quick run for automated notebook validation
dry_run = False
```

--------------------------------

### Neural ODE Model Setup

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4g_gde_node_classification.ipynb

Demonstrates how to construct a Neural ODE model using a sequential model that incorporates GCN layers and the NeuralODE solver.

```APIDOC
## NeuralODE Model Construction

### Description
Sets up a Neural ODE model by defining the ODE function using GCN layers and specifying the ODE solver and time span.

### Method
`NeuralODE(func, solver, s_span)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
func = nn.Sequential(
    GCNLayer(g=g, in_feats=64, out_feats=64, activation=nn.Softplus(), dropout=0.9),
    GCNLayer(g=g, in_feats=64, out_feats=64, activation=None, dropout=0.9)
).to(device)

neuralDE = NeuralODE(func, solver='rk4', s_span=torch.linspace(0, 1, 3)).to(device)
```

### Response
#### Success Response (200)
None

#### Response Example
None
```

--------------------------------

### Import Libraries and Setup

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_odeint.ipynb

Imports necessary libraries for numerical computation, neural networks, and plotting. It also sets up autoreload for development.

```python
import torchdyn
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torchdyn.numerics import Euler, RungeKutta4, Tsitouras45, DormandPrince45
from torchdyn.numerics import odeint
from torchdyn.core import ODEProblem

import torchdiffeq
import time 
%load_ext autoreload
%autoreload 2
```

--------------------------------

### Setup and Imports for Neural SDEs

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module1-neuralde/m1e_neural_sde_cookbook.ipynb

Imports necessary libraries like torch and PyTorch Lightning, and sets up the execution environment. It also configures the device for computation (GPU or CPU).

```python
import torch

%load_ext autoreload
%autoreload 2
```

```python
import sys
sys.path.append('..')
```

```python
from torchdyn.utils import plot_3D_dataset
from torchdyn.datasets import ToyDataset
from torch.utils.data import TensorDataset, DataLoader
```

```python
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
dry_run = False
```

--------------------------------

### Setup Controlled System and Policy Networks (Python)

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3b_optimal_control.ipynb

Initializes the prior and target distributions, defines neural networks for the energy shaping and damping injection policies (V and K), and sets up the controlled system dynamics. This prepares the system for training.

```python
from numpy import pi as pi
import torch
import torch.nn as nn
from torchdyn.systems import ControlledSystem, AugmentedDynamics

prior = prior_dist(-2*pi, 2*pi, -2*pi, 2*pi) # Uniform "prior" distribution of initial conditions x(0) 
target = target_dist([0, 0], [.001, .001]) # Normal target distribution for x(T)

# define optimal energy shaping policy networks
hdim = 64
V = nn.Sequential(
          nn.Linear(1, hdim),
          nn.Softplus(), 
          nn.Linear(hdim, hdim),
          nn.Tanh(), 
          nn.Linear(hdim, 1))
K = nn.Sequential(
          nn.Linear(2, hdim),
          nn.Softplus(),
          nn.Linear(hdim, 1),
          nn.Softplus())

# init to zero par.s of the final layer
for p in V[-1].parameters(): torch.nn.init.zeros_(p)
for p in K[-2].parameters(): torch.nn.init.zeros_(p)

# define controlled system dynamics
f = ControlledSystem(V, K)
aug_f = AugmentedDynamics(f, ControlEffort(f))

# define time horizon
t_span = torch.linspace(0, 3, 30)
```

--------------------------------

### Import Libraries and Setup

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_symplectic.ipynb

Imports necessary libraries for PyTorch, torchdyn, and plotting. It also loads the autoreload extension for development.

```python
import torchdyn
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torchdyn.numerics import AsynchronousLeapfrog
from torchdyn.numerics import odeint, odeint_symplectic

import torchdiffeq
import time 
%load_ext autoreload
%autoreload 2
```

--------------------------------

### Initialize Neural ODE Training Environment

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module1-neuralde/m1a_neural_ode_cookbook.ipynb

Imports necessary torchdyn modules and configures the environment for training. This setup includes loading autoreload extensions and defining the device for computation.

```python
from torchdyn.core import NeuralODE
from torchdyn.nn import DataControl, DepthCat, Augmenter, GalLinear, Fourier
from torchdyn.datasets import *
from torchdyn.utils import *
%load_ext autoreload
%autoreload 2
dry_run = False
```

--------------------------------

### Initialize Stochastic Hybrid System Environment

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2c_hybrid_odeint.ipynb

Imports necessary libraries including torch, torchdyn, and matplotlib to prepare for hybrid system simulation. This setup is required to utilize the odeint_hybrid solver for stochastic event handling.

```python
import torch
import attr
from torch.distributions import Exponential
import torch.nn as nn
import matplotlib.pyplot as plt
from torchdyn.numerics import odeint_hybrid

%load_ext autoreload
%autoreload 2
```

--------------------------------

### Import Libraries and Initialize System for ODE Solving

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2b_multiple_shooting.ipynb

Imports necessary libraries from PyTorch, TorchDyn, and Matplotlib. It also initializes the Lorenz system and defines the time span for integration. This setup is crucial for running ODE solvers.

```python
import torchdyn
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torchdyn.numerics import Euler, RungeKutta4, Tsitouras45, DormandPrince45, MSZero, MSBackward
from torchdyn.numerics import odeint, odeint_mshooting, Lorenz

import torchdiffeq
import time 
%load_ext autoreload
%autoreload 2
```

```python
# quick run for automated notebook validation
dry_run = False
```

```python
x0 = torch.randn(8, 3) + 15
t_span = torch.linspace(0, 3, 3000)
sys = Lorenz()
```

--------------------------------

### Import Libraries and Setup

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2a_hypersolver_odeint.ipynb

Imports necessary libraries like torch, torch.nn, and matplotlib, along with specific modules from torchdyn for ODE solving and datasets. It also includes notebook magic commands for autoreload.

```python
import time
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

from torchdyn.core import NeuralODE
from torchdyn.datasets import *
from torchdyn.numerics import odeint, Euler, HyperEuler

%load_ext autoreload
%autoreload 2
```

--------------------------------

### Setup Optimal Control with Neural Controller in Python

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4b_hypersolver_optimal_control.ipynb

Reinstantiates a neural network model to act as a controller for a system and sets up the time span and initial distribution for optimal control problems. This is a preparatory step for training the neural controller.

```python
# Reinstantiate controller
model = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 1)).to(device)
u = NeuralController(model) 
for p in u.model[-1].parameters(): torch.nn.init.zeros_(p)
sys.u = u

# Time span
t0, tf = 0, 2 # initial and final time for controlling the system
steps = 20 + 1 # so we have a time step of 0.1s
t_span = torch.linspace(t0, tf, steps).to(device)

# Initial distribution
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
```

--------------------------------

### Install Torchdyn via pip

Source: https://github.com/diffeqml/torchdyn/blob/master/docs/index.rst

This command installs the torchdyn library and its dependencies using the Python package manager. It is the primary method for setting up the environment to use continuous neural architectures.

```bash
pip install torchdyn
```

--------------------------------

### Initialize Lorenz System and Time Span in Python

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4c.multiple_shooting_layers.ipynb

Initializes the Lorenz system and defines the time span for integration. This setup is common for demonstrating ODE solvers.

```python
import torchdyn
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torchdyn.core import MultipleShootingLayer, MultipleShootingProblem
from torchdyn.numerics import Lorenz

import torchdiffeq
import time 
%load_ext autoreload
%autoreload 2

x0 = torch.randn(8, 3) + 15
t_span = torch.linspace(0, 3, 3000)
sys = Lorenz()
```

--------------------------------

### Initialize and Solve ODE Systems

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_stiff.ipynb

Demonstrates the initialization of an ODEProblem using torchdyn and compares trajectory computation against torchdiffeq's odeint_adjoint. It measures execution time for both solvers to evaluate performance differences.

```python
prob = ODEProblem(f, sensitivity='interpolated_adjoint', solver='dopri5', atol=1e-3, rtol=1e-3, atol_adjoint=1e-3, rtol_adjoint=1e-3)
t_eval, sol_torchdyn = prob.odeint(x, t_span)
t0 = time.time()
sol_torchdiffeq = torchdiffeq.odeint_adjoint(f, x, t_span, method='dopri5', atol=1e-3, rtol=1e-3)
t_end2 = time.time() - t0
true_sol = torchdiffeq.odeint_adjoint(f, x, t_span, method='dopri5', atol=1e-9, rtol=1e-9)
```

--------------------------------

### Training Neural ODEs with PyTorch Lightning

Source: https://context7.com/diffeqml/torchdyn/llms.txt

Provides a comprehensive example of training a Neural ODE for classification tasks using the PyTorch Lightning framework.

```APIDOC
## Training Neural ODEs with PyTorch Lightning

### Description
This section demonstrates how to integrate Torchdyn's `NeuralODE` module with PyTorch Lightning for streamlined training and experimentation. It covers dataset generation, model definition, and the PyTorch Lightning training loop.

### Method
N/A (Illustrative code)

### Endpoint
N/A

### Parameters
N/A

### Request Example
```python
import torch
import torch.nn as nn
import torch.utils.data as data
import pytorch_lightning as pl
from torchdyn.core import NeuralODE
from torchdyn.datasets import ToyDataset

# Generate dataset
d = ToyDataset()
X, y = d.generate(n_samples=1024, dataset_type='moons', noise=0.1)

train_dataset = data.TensorDataset(X, y.long())
train_loader = data.DataLoader(train_dataset, batch_size=128, shuffle=True)

# Define NeuralODE model
vector_field = nn.Sequential(
    nn.Linear(2, 64),
    nn.Tanh(),
    nn.Linear(64, 64),
    nn.Tanh(),
    nn.Linear(64, 2)
)

neural_ode = NeuralODE(
    vector_field,
    sensitivity='adjoint',
    solver='dopri5',
    atol=1e-3,
    rtol=1e-3
)

# Example PyTorch Lightning Module (simplified)
class ODEModule(pl.LightningModule):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, x):
        # Assuming x is a tuple (t_span, initial_state)
        t_span, x0 = x
        return self.model(t_span, x0)

    def training_step(self, batch, batch_idx):
        x, y = batch
        # Assuming x is the initial state and we need to define t_span
        t_span = torch.linspace(0, 1, 10) # Example time span
        pred = self.model(t_span, x)
        loss = nn.CrossEntropyLoss()(pred, y)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

# Instantiate and train
model = ODEModule(neural_ode)
# trainer = pl.Trainer(max_epochs=10)
# trainer.fit(model, train_loader)
```

### Response
N/A
```

--------------------------------

### Set Device for PyTorch

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3a_image_classification.ipynb

Determines whether to use a CUDA-enabled GPU or the CPU for computations. This is a common setup for deep learning tasks to leverage hardware acceleration.

```python
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
```

--------------------------------

### Initialize Device for Training

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4f_stable_neural_odes.ipynb

Sets up the computation device (CPU or GPU) for training the neural network models. This ensures that the computations are performed on the available hardware, optimizing performance.

```python
import torch.utils.data as data
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
```

--------------------------------

### Initialize ODE Solvers and Problem Definitions

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/00_quickstart.ipynb

Imports necessary numerical solvers and problem wrappers from the torchdyn library for setting up differential equation simulations.

```python
import time
from torchdyn.numerics import Euler, RungeKutta4, Tsitouras45, DormandPrince45, MSZero, Euler, HyperEuler
from torchdyn.numerics import odeint, odeint_mshooting, Lorenz
from torchdyn.core import ODEProblem, MultipleShootingProblem
```

--------------------------------

### Setting up PyTorch and Lightning Dependencies

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3a_image_classification.ipynb

Initializes the necessary PyTorch and PyTorch Lightning modules for data loading, model training, and logging metrics during the training process.

```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.metrics.functional import accuracy
```

--------------------------------

### Initialize Neural ODE for CNF

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3c_continuous_normalizing_flows.ipynb

Configures a Continuous Normalizing Flow (CNF) using a Neural ODE. This setup includes the trace estimator for divergence calculation and solver parameters for the ODE integration.

```python
cnf = CNF(f, trace_estimator=autograd_trace)
nde = NeuralODE(cnf, solver='dopri5', sensitivity='adjoint', atol=1e-4, rtol=1e-4)
```

--------------------------------

### Initialize a NeuralODE Model

Source: https://github.com/diffeqml/torchdyn/blob/master/README.md

Demonstrates how to define a neural network architecture using PyTorch modules and wrap it within a NeuralODE class to create a trainable differential equation model.

```python
from torchdyn.core import NeuralODE
import torch.nn as nn

# Define the underlying neural network
f = nn.Sequential(nn.Conv2d(1, 32, 3),
                  nn.Softplus(),
                  nn.Conv2d(32, 1, 3)
          )

# Initialize the NeuralODE with the network
nde = NeuralODE(f)
```

--------------------------------

### Depth-Invariant Conv Neural ODE

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3a_image_classification.ipynb

This section demonstrates how to implement a depth-invariant convolutional Neural ODE model. It includes the model definition using `NeuralODE` and `Augmenter`, followed by the training setup using `Learner` and `pl.Trainer`.

```APIDOC
## Depth-Invariant Conv Neural ODE

### Description
This example shows how to define and train a depth-invariant convolutional Neural ODE model using TorchDyn. The model incorporates `Augmenter` for data augmentation and `NeuralODE` for the ODE solver.

### Method
**Model Definition**: Uses `torch.nn.Sequential` to combine `Augmenter`, `NeuralODE`, convolutional layers, and a linear layer.
**Training**: Utilizes `Learner` and `pytorch_lightning.Trainer` for the training loop.

### Parameters
**NeuralODE Solver**: `rk4`
**NeuralODE Sensitivity**: `autograd`
**Augmenter Dimensions**: `augment_dims=10`
**Trainer**: `max_epochs=3`, `progress_bar_refresh_rate=1`, `gpus=1`

### Request Example
```python
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torchdyn.core import NeuralODE
from torchdyn.datasets import DatedSync, ECG, Energy, Gas, Gravity, Protein, Sine
from torchdyn.models import Augmenter
from torchdyn.training import Trainer, Learner

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Define the vector field function
func = nn.Sequential(nn.Conv2d(11, 11, 3, padding=1),
                     nn.Tanh(),                 
                     ).to(device)

# Initialize NeuralODE
neuralDE = NeuralODE(func, 
                     solver='rk4',
                     sensitivity='autograd').to(device)

# Define the full model
model = nn.Sequential(Augmenter(augment_dims=10),
                      neuralDE,
                      nn.Conv2d(11, 1, 3, padding=1),
                      nn.Flatten(),                     
                      nn.Linear(28*28, 10)).to(device)

# Initialize Learner and Trainer
learn = Learner(model)
trainer = pl.Trainer(max_epochs=3,
                     progress_bar_refresh_rate=1,
                     gpus=1
                     )

# Train the model
trainer.fit(learn)
```

### Response
#### Success Response (200)
Training output from PyTorch Lightning, including progress bars and model summary.

#### Response Example
```
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name  | Type       | Params
-------------------------------------
0 | model | Sequential | 9.1 K 
-------------------------------------
9.1 K     Trainable params
0         Non-trainable params
9.1 K     Total params
0.036     Total estimated model params size (MB)

```

### Error Handling
Potential errors include tensor shape mismatches during convolution operations, indicated by `RuntimeError` exceptions.

```

--------------------------------

### Define Loss Functions and Control Effort (Python)

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3b_optimal_control.ipynb

Implements a weighted negative log-likelihood loss function and a ControlEffort module to calculate the integral cost of control actions. These are used during the training process to guide the policy optimization.

```python
import torch.nn as nn
import torch

def weighted_log_likelihood_loss(x, target, weight):
    # weighted negative log likelihood loss
    log_prob = target.log_prob(x)
    weighted_log_p = weight * log_prob
    return -torch.mean(weighted_log_p.sum(1))

class ControlEffort(nn.Module):
    # control effort integral cost
    def __init__(self, f):
        super().__init__()
        self.f = f
    def forward(self, t, x):
        with torch.set_grad_enabled(True):
            q = x[:,:1].requires_grad_(True)
            u = self.f._energy_shaping(q) + self.f._damping_injection(x)
        return torch.abs(u)
```

--------------------------------

### Initialize and Train Neural ODE

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module1-neuralde/m1a_neural_ode_cookbook.ipynb

Demonstrates how to define a NeuralODE model with specific solvers and adjoint methods, then train it using a PyTorch Lightning Learner.

```python
model = NeuralODE(f, sensitivity='interpolated_adjoint', solver='tsit5', atol=1e-3, rtol=1e-3).to(device)
learn = Learner(t_span, model)
if dry_run: trainer = pl.Trainer(min_epochs=1, max_epochs=1)
else: trainer = pl.Trainer(min_epochs=150, max_epochs=250)
trainer.fit(learn)
```

--------------------------------

### Instantiate NeuralODE Model (Python)

Source: https://github.com/diffeqml/torchdyn/blob/master/docs/tutorials/quickstart.ipynb

Instantiates the `NeuralODE` class from `torchdyn`. It takes the defined vector field `f`, sensitivity method ('adjoint'), and ODE solver ('dopri5') as arguments. The model is then moved to the specified device.

```python
model = NeuralODE(f, sensitivity='adjoint', solver='dopri5').to(device)
```

--------------------------------

### Initialize and Train the Model

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4a_approximate_continuous_normalizing_flows.ipynb

Instantiates the `Learner` with the defined model and sets up a PyTorch Lightning `Trainer`. The `trainer.fit(learn)` command initiates the training process for a specified number of epochs. The output shows the model's parameter count and training progress.

```python
learn = Learner(model)
trainer = pl.Trainer(max_epochs=600)
trainer.fit(learn);
```

--------------------------------

### Training Neural ODEs with PyTorch Lightning Integration

Source: https://context7.com/diffeqml/torchdyn/llms.txt

This snippet demonstrates how to train a NeuralODE model for classification using PyTorch Lightning. It includes dataset generation, DataLoader setup, NeuralODE model definition with specified solver and tolerances, and integration into a PyTorch Lightning module.

```python
import torch
import torch.nn as nn
import torch.utils.data as data
import pytorch_lightning as pl
from torchdyn.core import NeuralODE
from torchdyn.datasets import ToyDataset

# Generate dataset
d = ToyDataset()
X, y = d.generate(n_samples=1024, dataset_type='moons', noise=0.1)

train_dataset = data.TensorDataset(X, y.long())
train_loader = data.DataLoader(train_dataset, batch_size=128, shuffle=True)

# Define NeuralODE model
vector_field = nn.Sequential(
    nn.Linear(2, 64),
    nn.Tanh(),
    nn.Linear(64, 64),
    nn.Tanh(),
    nn.Linear(64, 2)
)

neural_ode = NeuralODE(
    vector_field,
    sensitivity='adjoint',
    solver='dopri5',
    atol=1e-3,
    rtol=1e-3
)
```

--------------------------------

### Setting up ODE Problem and Initial Conditions

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_stiff.ipynb

Initializes the Van der Pol system with a specific alpha value and sets up the initial state `x` and time span `t_span` for the ODE integration.

```python
f = VanDerPol(20)
x = torch.randn(1024, 2)
t_span = torch.linspace(0, 1, 500)
```

--------------------------------

### Define Neural ODE Vector Field (Python)

Source: https://github.com/diffeqml/torchdyn/blob/master/docs/tutorials/quickstart.ipynb

Defines the vector field $f(h,\theta)$ for a Neural ODE using a PyTorch nn.Module. This example uses a simple MLP with a Tanh activation function. The module takes state `x` and time `t` as input.

```python
f = nn.Sequential(
        nn.Linear(2, 16),
        nn.Tanh(),
        nn.Linear(16, 2)                       
    )
t_span = torch.linspace(0, 1, 5)
```

--------------------------------

### Define Neural ODE Model with Higher-Order Augmentation

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module1-neuralde/m1c_augmentation_strategies.ipynb

This example demonstrates defining a Neural ODE model that utilizes higher-order augmentation. It specifies the ODE function, the NeuralODE layer with a specific solver ('dopri5') and order (2), and the model architecture including an Augmenter layer.

```python
func = nn.Sequential(nn.Linear(6, 64),
                     nn.Tanh(),
                     nn.Linear(64, 3))

neuralDE = NeuralODE(func, solver='dopri5', order=2).to(device)

model = nn.Sequential(Augmenter(augment_dims=3),
                      neuralDE,
                      nn.Linear(6,2)).to(device)
```

--------------------------------

### Initialize ODE Problem for Adjoint Sensitivity

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_odeint.ipynb

Sets up an ODEProblem in torchdyn with the 'adjoint' sensitivity method and Dopri5 solver. It also initializes the input tensor with requires_grad=True for gradient calculations.

```python
from torchdyn.core import ODEProblem

f = VanDerPol(0.5)
t_span = torch.linspace(0, 2, 300)
x = torch.randn(1024, 2, requires_grad=True)
```

--------------------------------

### Solving ODEs with Multiple Solvers in TorchDyn

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2a_hypersolver_odeint.ipynb

This snippet demonstrates solving an ODE system with different initial conditions using various ODE solvers provided by TorchDyn. It compares the computational time and accuracy of a custom hypersolver against standard numerical methods like Euler, RK4, and Dopri5. The code initializes random initial conditions, moves them to the appropriate device, and then calls `odeint` with different solver configurations.

```python
x0 = torch.randn(5, 3) + 14
hypersolver = hypersolver.to(device)
x0 = x0.to(device)

t0 = time.time()
t_eval, sol = odeint(sys, x0, t_span, solver=hypersolver)
sol = sol.detach().cpu()
hyper_sol_time = time.time() - t0

t0 = time.time()
t_eval, base_sol = odeint(sys, x0, t_span, solver='euler')
base_sol = base_sol.detach().cpu()
base_sol_time = time.time() - t0

t0 = time.time()
t_eval, rk4_sol = odeint(sys, x0, t_span, solver='rk4')
rk4_sol = rk4_sol.detach().cpu()
rk4_sol_time = time.time() - t0

t0 = time.time()
t_eval, dp5_low_sol = odeint(sys, x0, t_span, solver='dopri5', atol=1e-3, rtol=1e-3)
dp5_low_sol = dp5_low_sol.detach().cpu()
dp5_low_time = time.time() - t0
```

--------------------------------

### Solve ODE with TorchDyn using odeint_adjoint

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/gradcheck_tspan.ipynb

This example shows how to solve an ODE using `torchdiffeq.odeint_adjoint`. It defines a `VectorField` class to wrap the underlying function `f`, initializes parameters, and then iteratively solves the ODE, calculates a loss, and optimizes the time span `T`. This method is suitable for models requiring adjoint differentiation.

```python
class VectorField(nn.Module):
    def __init__(self, f):
        super().__init__()
        self.f = f
    def forward(self, t, x):
        return self.f(x)

sys = VectorField(f)
x = torch.zeros(1, 1, requires_grad=True) + 0.5
t0 = torch.zeros(1)
T = torch.ones(1).requires_grad_(True)
opt = torch.optim.Adam((T,), lr=1e-2)

for i in range(1000):
    t_span = torch.cat([t0, T])
    traj = torchdiffeq.odeint_adjoint(sys, x, t_span, method='dopri5', atol=1e-4, rtol=1e-4)
    loss = ((traj[-1] - torch.tensor([2]))**2).mean()
    print(f'L: {loss.item():.2f}, T: {t_span[-1].item():.2f}, xT: {traj[-1].item():.2f}', end='\r')
    loss.backward(); opt.step(); opt.zero_grad()
```

--------------------------------

### Initialize System and Data

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_symplectic.ipynb

Initializes the `TestSystem` with the defined vector field and creates random input data `x` and a time span `t_span` for integration.

```python
f = TestSystem(vf)
x = torch.randn(1024, 2)
t_span = torch.linspace(0, 4, 200)
```

--------------------------------

### Initialize State Distribution with PyTorch

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4b_hypersolver_optimal_control.ipynb

Sets up a uniform distribution for the initial state of a system. The state is defined in radians and radians per second, with limits specified by x0.

```python
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
```

--------------------------------

### Configure Sensitivity Analysis with ODEProblem

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/00_quickstart.ipynb

Wraps a system in an ODEProblem to perform sensitivity analysis using adjoint methods for optimization.

```python
prob = ODEProblem(sys, sensitivity='interpolated_adjoint', solver='dopri5', atol=1e-3, rtol=1e-3, solver_adjoint='tsit5', atol_adjoint=1e-3, rtol_adjoint=1e-3)
t0 = time.time()
t_eval, sol_torchdyn = prob.odeint(x0, t_span)
t_end1 = time.time() - t0
```

--------------------------------

### Basic Usage of Hypersolver API

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module2-numerics/m2a_hypersolver_odeint.ipynb

This section demonstrates the basic usage of the Hypersolver API in Torchdyn, showcasing how hypersolvers are integrated into the `odeint` API. It includes setting up an ODE system, solving it with a standard solver, and then using a hypersolver for comparison.

```APIDOC
## Basic Usage of the `Hypersolver` API

This section showcases the `torchdyn` API for [hypersolvers](https://arxiv.org/pdf/2007.09601.pdf). Hypersolvers, or hybrid ODE solvers equipped with a neural network to approximate residuals, are seamlessly integrated into the `odeint` API. A major design of the API is exactly preserving a persistent state for the solver, which in the case of standard ODE solvers contains the Tableau. For hypersolvers, the state also contains their hypernetwork parameters.

### Setup and Standard ODE Solution

This code block sets up the necessary imports, defines a Lorenz system, and solves it using the `tsit5` solver.

```python
import time
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

from torchdyn.core import NeuralODE
from torchdyn.datasets import *
from torchdyn.numerics import odeint, Euler, HyperEuler

%load_ext autoreload
%autoreload 2
```

```python
# quick run for automated notebook validation
dry_run = False
```

```python
class Lorenz(nn.Module):
    def __init__(self):
        super().__init__()
        self.p = nn.Linear(1,1)
    
    def forward(self, t, x):
        x1, x2, x3 = x[...,:1], x[...,1:2], x[...,2:]
        dx1 = 10 * (x2 - x1)
        dx2 = x1 * (28 - x3) - x2
        dx3 = x1 * x2 - 8/3 * x3
        return torch.cat([dx1, dx2, dx3], -1)
```

```python
x0 = torch.randn(512, 3) + 15 # solve 512 IVPs in parallel!
t_span = torch.linspace(0, 5, 4000)
sys = Lorenz()
if dry_run: t_eval, sol_gt = odeint(sys, x0, t_span, solver='tsit5', atol=1e-3, rtol=1e-3)
else: t_eval, sol_gt = odeint(sys, x0, t_span, solver='tsit5', atol=1e-8, rtol=1e-8)
```

### Visualization of Ground Truth Solution

This code block visualizes the first few trajectories of the ground truth solution obtained from the `tsit5` solver.

```python
fig = plt.figure(figsize=(15, 4))
axs = fig.subplots(3,1)
axs[0].plot(sol_gt[:,:4,0], c='black'); axs[1].plot(sol_gt[:,:4,1], c='r'); axs[2].plot(sol_gt[:,:4,2], c='b');
```

```
Result:
<Figure size 1080x288 with 3 Axes>
```

### Introducing and Using a Hypersolver

This section introduces a `VanillaHyperNet` and uses `HyperEuler` to solve the ODE system, demonstrating potential speed-up.

```python
class VanillaHyperNet(nn.Module):
    def __init__(self, net):
        super().__init__()
        self.net = net
        for p in self.net.parameters():
            torch.nn.init.uniform_(p, 0, 1e-5)
    def forward(self, t, x):
        return self.net(x)
```

```python
net = nn.Sequential(nn.Linear(3, 64), nn.Softplus(), nn.Linear(64, 64), nn.Softplus(), nn.Linear(64, 3))
hypersolver = HyperEuler(VanillaHyperNet(net))
t_eval, sol = odeint(sys, x0, t_span, solver=hypersolver)
sol = sol.detach()
```

### Visualizing Hypersolver Solution and Errors

This code block visualizes the hypersolver's solution and compares it with the ground truth, highlighting the errors.

```python
fig = plt.figure(figsize=(15, 14))
axs = fig.subplots(6,1)
axs[0].plot(sol[:,:4,0], c='black'); axs[2].plot(sol[:,:4,1], c='r'); axs[4].plot(sol[:,:4,2], c='b');
axs[0].set_xticks([]); axs[2].set_xticks([]); axs[4].set_xticks([])
# The error accumulates as Euler (base solver of HyperEuler) does not converge here. 
# Here we plot the errors (mean across batches of initial conditions)
axs[1].plot((sol-sol_gt).abs().mean(1).sum(1), c='black')
axs[1].set_title('Error on $x_0$')
axs[3].plot((sol-sol_gt).abs().mean(1).sum(1), c='r')
axs[3].set_title('Error on $x_1$')
axs[5].plot((sol-sol_gt).abs().mean(1).sum(1), c='b');
axs[5].set_title('Error on $x_2$')
```

```
Result:
Text(0.5, 1.0, 'Error on $x_2$')
```

```
Result:
<Figure size 1080x1008 with 6 Axes>
```
```

```APIDOC
### Training a Hypersolver

This section details how to train a hypersolver by minimizing the residuals between the ground-truth solution and the base solver's prediction.

```python
base_solver = Euler()
X = sol_gt[:-1].reshape(-1, 3)
X_next_gt = sol_gt[1:].reshape(-1, 3)
# step forward (fixed-step, time-invariant system hence any `t` as first argument is fine) with base solver
dt = t_span[1] - t_span[0]
_, X_next, _ = base_solver.step(sys, X, 0., dt) # step returns a Tuple (k1, berr, sol). The first two values are used internally
# within `odeint`
residuals = (X_next_gt - X_next) / dt**2
```

```python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
model = nn.DataParallel(hypersolver, device_ids=[0,1]) # feel free to change here according to your setup and GPU available.
model = model.to(device)
X = X.to(device)
residuals = residuals.to(device)
```

```python
# training will take a while...
opt = torch.optim.Adadelta(model.parameters(), lr=3e-4)
if dry_run: STEPS = 2
else: STEPS = 300000


loss_func = nn.MSELoss()
hypernet = model.module.hypernet
for k in range(STEPS):
    residuals_hypersolver = hypernet(0., X)
    loss = loss_func(residuals, residuals_hypersolver)
    print(f'Step: {k}, Residual loss: {loss:.3f}', end='\r')
    loss.backward(); opt.step(); opt.zero_grad()
```

```
Output:
Step: 299999, Residual loss: 134.91475
```

### Performance Comparison

This section briefly discusses the performance comparison, noting that the hypersolver can be roughly twice as fast as a `rk4` solver with significantly higher accuracy.

```
# roughly twice as fast as a rk4 (3 orders higher). How does the accuracy compare to the base solver (and others)?
```
```

--------------------------------

### Prepare dataset and dataloader

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4e_lagrangian_nets.ipynb

Creates a synthetic dataset for training and a PyTorch DataLoader to iterate over it. The dataset consists of initial conditions (X) and corresponding second derivatives (Xdd), suitable for training a dynamics model.

```python
import torch.utils.data as data
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

m, k, l = 1, 1, 1
X = torch.Tensor(2**14, 2).uniform_(-1, 1).to(device)
Xdd = -k*X[:,0]/m

train = data.TensorDataset(X, Xdd)
trainloader = data.DataLoader(train, batch_size=64, shuffle=False)
```

--------------------------------

### Data Preparation for Neural SDE Training

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module1-neuralde/m1e_neural_sde_cookbook.ipynb

Generates a toy dataset using `ToyDataset` and prepares it for training. This includes creating tensors, moving them to the appropriate device, and setting up a DataLoader.

```python
d = ToyDataset()
X, yn = d.generate(n_samples=512, dataset_type='moons', noise=.4)
X_train = torch.Tensor(X).to(device)
y_train = torch.LongTensor(yn.long()).to(device)
train = TensorDataset(X_train, y_train)
trainloader = DataLoader(train, batch_size=len(X), shuffle=False)
```

--------------------------------

### Load and Visualize Toy Dataset

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module4-model/m4a_approximate_continuous_normalizing_flows.ipynb

Loads a toy dataset ('diffeqml' logo) and visualizes it using Matplotlib. The data is preprocessed by centering and scaling it. This step is crucial for understanding the data distribution before applying density estimation.

```python
data = ToyDataset()
n_samples = 1 << 14
n_gaussians = 7

X, yn = data.generate(n_samples, 'diffeqml', noise=5e-2)
X = (X - X.mean())/X.std()

import matplotlib.pyplot as plt
plt.figure(figsize=(3, 3))
plt.scatter(X[:,0], X[:,1], c='olive', alpha=0.3, s=1)
```

--------------------------------

### Comparing TorchDyn and TorchDiffEq Solvers (Adaptive)

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_stiff.ipynb

Compares torchdyn.odeint_mshooting with torchdiffeq.odeint using an implicit Euler solver for torchdyn and dopri5 for torchdiffeq. It measures and prints the execution times. A high-accuracy dopri5 solution is used as a reference.

```python
t0 = time.time()
t_eval, sol1 = odeint_mshooting(f, x, t_span, solver='ieuler')
t_end1 = time.time() - t0
print(t_end1)

t0 = time.time()
sol2 = torchdiffeq.odeint(f, x, t_span, method='dopri5', atol=1e-4, rtol=1e-4)
t_end2 = time.time() - t0
print(t_end2)

true_sol = torchdiffeq.odeint(f, x, t_span, method='dopri5', atol=1e-9, rtol=1e-9)
```

--------------------------------

### Initialize Torchdyn Environment

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/00_quickstart.ipynb

Imports necessary modules from torchdyn and configures the notebook environment for development.

```python
from torchdyn.core import NeuralODE
from torchdyn.datasets import *
from torchdyn import *

%load_ext autoreload
%autoreload 2
dry_run = False
```

--------------------------------

### Simulate and Visualize Trajectories (Python)

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module3-tasks/m3b_optimal_control.ipynb

Generates simulation trajectories using odeint from torchdiffeq and visualizes various system properties like control inputs (u_es, u_di), total control effort (u), and energy functions (H, Hd, Ha) on a grid. This helps in analyzing the learned policy.

```python
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from torchdiffeq import odeint

n_ic = 256
x0 = prior.sample(torch.Size([n_ic])).cpu()
x0 = torch.cat([x0, torch.zeros(n_ic, 1)], 1)
model = aug_f.cpu()
#
traj = odeint(model, x0, t_span, method='midpoint').detach()
traj = traj[..., :-1]
u_es = torch.cat([model.f._energy_shaping(q.requires_grad_()) for q in traj[:,:,:1]], 1).detach()
u_di = torch.cat([model.f._damping_injection(x) for x in traj], 1).detach()
u = (u_es + u_di)
H = torch.cat([model.f._autonomous_energy(x) for x in traj], 1).detach()
Hd = torch.cat([model.f._energy(x) for x in traj], 1).detach()
Ha = Hd - H
# plots on grid
n_grid = 100
q_grid, p_grid = torch.linspace(-2*pi, 2*pi, n_grid), torch.linspace(-2*pi, 2*pi, n_grid)
Q, P = torch.meshgrid(q_grid, p_grid) ; z = torch.cat([Q.reshape(-1, 1), P.reshape(-1, 1)], 1)
H_grid = model.f._autonomous_energy(z).reshape(n_grid, n_grid).detach()
Hd_grid = model.f._energy(z).reshape(n_grid, n_grid).detach()
Ha_grid = Hd_grid - H_grid
Kd_grid = -model.f._damping_injection(z).reshape(n_grid, n_grid).detach()/P   
```

--------------------------------

### Initialize NeuralSDE Model

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/module1-neuralde/m1e_neural_sde_cookbook.ipynb

Initializes the `NeuralSDE` model from `torchdyn.core`. This involves defining the drift (`f`) and diffusion (`g`) functions, specifying solver parameters, and setting up the Brownian motion.

```python
from torchdyn.core import NeuralSDE
import torchsde

t_span = torch.linspace(0, 0.1, 100).to(device)
size = X.shape

bm = torchsde.BrownianInterval(
    t0=t_span[0],
    t1=t_span[-1],
    size=size,
    device=device,
    levy_area_approximation='space-time'
)
```

```python
f = nn.Sequential(nn.Linear(2, 64), nn.Tanh(), nn.Linear(64, 2))
g = nn.Sequential(nn.Linear(2, 64), nn.Tanh(), nn.Linear(64, 2))


model = NeuralSDE(f, 
                  g,
                  solver='euler',
                  noise_type='diagonal',
                  sde_type='ito',
                  sensitivity='autograd',
                  t_span=t_span,
                  bm=bm
                 ).to(device)
```

--------------------------------

### Solve ODE Systems with Functional API

Source: https://github.com/diffeqml/torchdyn/blob/master/tutorials/00_quickstart.ipynb

Demonstrates how to define a Lorenz system and solve it using various numerical solvers like dopri5, euler, and rk4, measuring execution time for each.

```python
x0 = torch.randn(8, 3) + 15
t_span = torch.linspace(0, 3, 2000)
sys = Lorenz()
t0 = time.time()
t_eval, accurate_sol = odeint(sys, x0, t_span, solver='dopri5', atol=1e-6, rtol=1e-6)
accurate_sol_time = time.time() - t0
t0 = time.time()
t_eval, base_sol = odeint(sys, x0, t_span, solver='euler')
base_sol_time = time.time() - t0
t0 = time.time()
t_eval, rk4_sol = odeint(sys, x0, t_span, solver='rk4')
rk4_sol_time = time.time() - t0
t0 = time.time()
t_eval, dp5_low_sol = odeint(sys, x0, t_span, solver='dopri5', atol=1e-3, rtol=1e-3)
dp5_low_time = time.time() - t0
t0 = time.time()
t_eval, ms_sol = odeint_mshooting(sys, x0, t_span, solver='mszero', fine_steps=2, maxiter=4)
ms_sol_time = time.time() - t0
```

--------------------------------

### Comparing Adjoint Sensitivity Analysis

Source: https://github.com/diffeqml/torchdyn/blob/master/benchmarks/numerics/bench_stiff.ipynb

Sets up an ODEProblem in TorchDyn with adjoint sensitivity enabled and compares its integration time and accuracy against torchdiffeq.odeint_adjoint. Both use the dopri5 solver with specified tolerances.

```python
from torchdyn.core import ODEProblem

f = VanDerPol(0.1)
t_span = torch.linspace(0, 3, 300)
x = torch.randn(1024, 2, requires_grad=True)


prob = ODEProblem(f, sensitivity='adjoint', solver='dopri5', atol=1e-3, rtol=1e-3, atol_adjoint=1e-3, rtol_adjoint=1e-3)
t0 = time.time()
t_eval, sol_torchdyn = prob.odeint(x, t_span)
t_end1 = time.time() - t0
print(t_end1)

t0 = time.time()
sol_torchdiffeq = torchdiffeq.odeint_adjoint(f, x, t_span, method='dopri5', atol=1e-3, rtol=1e-3)
t_end2 = time.time() - t0
print(t_end2)

true_sol = torchdiffeq.odeint_adjoint(f, x, t_span, method='dopri5', atol=1e-9, rtol=1e-9)
```