Quabla Developer Guide
This guide explains how to use Quabla and how it works inside. It is written
as a sequence of chapters: the first ones teach the Python API from the
ground up, the middle ones cover each library area in depth, and the last
ones describe the implementation for contributors.
The guide complements, and does not replace, the other documents:
- API Reference is the exhaustive statement of every function’s
contract, numerical convention, and backend behavior. The guide links to
it whenever a detail is too fine-grained for a tutorial.
- README is the project overview and the build matrix.
- jax_like_roadmap.md records per-feature status
and validation runs.
- CHANGELOG lists what changed in each release.
Every code example in the guide was run against Quabla v0.5.0 on the CPU
backend unless it says otherwise; printed output is shown in a text block
right after the code.
Chapters
Using Quabla
- Getting Started: installation, the first
gradient, and the mental model of eager arrays, traced values, and
transforms.
- Arrays and Operations: creating arrays, dtypes and
promotion, indexing, broadcasting, reductions, the math library, and NumPy
interoperability.
- Function Transforms:
grad, value_and_grad, jvp,
vjp, jacobian, hessian, vmap, and jit, and how they compose.
- Control Flow:
cond, fori_loop, scan, and
while_loop as differentiable IR regions.
- Custom Derivatives and Checkpointing:
custom_vjp, custom_jvp, checkpoint, and stop_gradient.
- Pytrees, Random Numbers, and Serialization:
structured parameters, keyed random sampling, and
save/load.
- Optimization: Adam, AdamW, SGD, learning-rate
schedules, gradient clipping, L-BFGS, and the retained-buffer
Trainer.
- Linear Algebra and Solvers:
quabla.linalg,
Krylov solvers, and Newton’s method with implicit gradients.
- Differential Equations:
quabla.ode.odeint with fixed-step,
adaptive, and stiff integrators.
- Devices and Compilation: running on CUDA and MLX,
precision, ahead-of-time lowering, the compiler facade, and data
parallelism.
- Tutorial: Training a PINN: an end-to-end
physics-informed neural network, from model to Adam-then-L-BFGS training
and device execution.
- Performance, Debugging, and Limitations:
tracing costs, error types, common pitfalls, and known limits.
Inside Quabla
- Architecture: crates, the tensor IR, autodiff
transforms, plan compilation, and the CPU, CUDA, and MLX backends.
- Contributing and Extending: the development
workflow, the test suites, and how to add a new operation.
- Migrating From the v0.1 API: mapping the v0.1
helpers that v1.0 removed to the core API.
Conventions
import quabla as qb is assumed in every example.- “The core API” means the JAX-style API introduced in v0.2 (
qb.grad,
qb.vmap, qb.jit, …). “The v0.1 API” means the spec-based helpers
(trace_tensor, tensor_*_fn) and the 2D Matrix API, which v1.0
removed.
- Paths such as
python/quabla/optim.py are relative to the repository
root.