Skip to content

Developer Setup and Workflow

Setup

git clone https://github.com/PREDICT-EPFL/anvil.git
cd anvil
uv sync

Testing

uv run pytest                    # All tests
uv run pytest tests/unit         # Unit tests only
uv run pytest tests/integration  # Integration tests (requires CMake build)

Code quality checks

uv run ruff check .   # Linting
uv run ruff format .  # Formatting
uv run ty check .     # Type checking

Project structure

src/anvil/
├── language.py            # Single point of entry for tinygrad (re-exports, sets defaults)
├── ad/                    # Automatic differentiation (jacobian, gradient, hessian, sparsity)
├── codegen/               # C++ code generation
│   ├── templates/         # Jinja2 templates for codegen
│   ├── common.py          # Shared utilities and base classes
│   └── numerical_function.py  # NumericalFunction and SparseNumericalFunction
├── transform/             # Graph transformations
│   ├── _vmap.py           # Graph-rewriting vmap (batch dimension insertion)
│   └── _jax_to_tiny.py    # JAX-to-tinygrad conversion
├── sqp/                   # SQP solver
│   ├── problem.py         # Problem and SparseProblem definitions
│   └── sqp_function.py    # SQPFunction codegen
├── ocp/                   # Optimal control problem support
│   ├── multistage.py      # MultistageFn and MultistageProblem
│   └── utils.py           # Low-level OCP sparse assembly helpers
└── ...

anvil.language

All tinygrad imports in anvil go through src/anvil/language.py. It re-exports types like Tensor, UOp, dtypes, Ops, etc. and sets two important defaults: - Default float type: float64 - Default device: CPU

When adding new tinygrad dependencies, add them to language.py and import from anvil.language in the rest of the codebase.