Adaptive Integration Time for Neural ODEs

Jax Implementation of the Adaptive Integration Time (AIT) algorithm for Neural ODEs.

Requirements

• Python 3.11
• uv for dependency management
• Optional: an NVIDIA GPU with CUDA 13 for the cuda extra

Installation

With uv (recommended). CPU-only:

uv sync

With CUDA 13 GPU support:

uv sync --extra cuda

This creates a virtual environment in .venv/ and installs the locked dependencies from uv.lock. Run commands with uv run, e.g.:

uv run python scripts/plot_experiments.py results/ait_mnist_0.001.csv results/node_mnist_0.csv

The experiment scripts below use uv run --no-sync, so they respect whichever extras you synced with (a plain uv run would re-sync without the cuda extra and uninstall the GPU wheels).

Alternatively, install into an existing environment with pip:

pip install -e .

Reproduce the experiments:

# Make the scripts executable
chmod +x experiments/run_ait.sh
chmod +x experiments/run_node.sh

# Run the experiments
./experiments/run_ait.sh g2
./experiments/run_node.sh g2

Citation

If you found this useful, please cite as:

@misc{grassoramos2026ait,
  title  = {Adaptive Integration Time for Neural ODEs},
  author = {Grasso Ramos, Lucas},
  year   = {2026},
  month  = aug,
  url    = {https://github.com/LucasGrasso/AIT}
}