Quabla

15. Migrating From the v0.1 API

Quabla v0.1 exposed transforms as spec-based helpers (trace_tensor, tensor_value_and_grad_fn, tensor_vmap_mlx_fn, …), with one spelling per combination of transform and backend. v0.2 replaced them with the composable core API. This chapter shows how to move existing code.

15.1 What v1.0 Removed

Through 0.x every v0.1 name kept working, and from v0.2 the names with a replacement emitted a DeprecationWarning on first access. v1.0 removes them all:

To find every use before upgrading, run your code on the last 0.x release (0.7.1) with the warnings turned into errors:

python -W error::DeprecationWarning your_script.py

15.2 The Conceptual Change

v0.1 Core API
Inputs are declared up front as (name, shape, dtype) specs Inputs are ordinary arguments; shapes and dtypes come from the values on the first call
Results are dicts keyed by input name Results mirror the argument pytrees
One helper per transform and backend (tensor_vmap_vjp_cuda_fn) Transforms compose (jit(vjp(vmap(f)), device="cuda"))
Derivatives of a trace via symbolic_jvp(name) on a TensorTraceResult grad, jvp, jacobian of a Python function
Bounded batch specialization via tensor_jit_batch_fn(max_specializations=4) Every transform caches up to max_traces signatures
Device optimizers (mlx_adam_loss_optimizer, cuda_adam_vjp_optimizer) optim.Trainer(..., device=...)

15.3 A Worked Example

The v0.1 version of examples/pinn_poisson.py fits one weight w so that sin(w x) solves u'' = -pi^2 sin(pi x) with zero boundary values. In v0.1 it builds two separately traced plans, takes symbolic JVPs by input name, and sums the gradients by hand:

residual_trace = quabla.trace_tensor(
    lambda x, weight, forcing: (x * weight).sin(),
    [("x", [5, 1]), ("weight", [1, 1]), ("forcing", [5, 1])],
)
second_derivative = residual_trace.symbolic_jvp("x").symbolic_jvp("x")
residual = second_derivative.output + second_derivative.graph.input("forcing")
residual_plan = (residual * residual).mean().compile_cpu()
# ... a second plan for the boundary loss ...
for _ in range(2000):
    _, residual_gradients = residual_plan.evaluate_value_and_vjp({...}, quabla.Tensor([], [1.0]))
    _, boundary_gradients = boundary_plan.evaluate_value_and_vjp({...}, quabla.Tensor([], [1.0]))
    parameters = optimizer.step(
        parameters,
        quabla.sum_gradients([residual_gradients, boundary_gradients], ["weight"]),
    )

With the core API, the model is written for one point, the loss is one function, and Trainer runs the loop:

import math
import quabla as qb

def u(x, w):
    return qb.sin(x * w)

u_xx = qb.vmap(qb.grad(qb.grad(u)), in_axes=(0, None))

def loss(params, x, x_boundary):
    residual = u_xx(x, params["weight"]) + math.pi**2 * qb.sin(math.pi * x)
    boundary = qb.vmap(u, in_axes=(0, None))(x_boundary, params["weight"])
    return qb.mean(residual ** 2) + qb.mean(boundary ** 2)

trainer = qb.optim.Trainer(
    loss, {"weight": qb.array(2.5)}, qb.optim.Adam(learning_rate=0.01),
    qb.array([0.15, 0.35, 0.55, 0.75, 0.9]), qb.array([0.0, 1.0]),
)
for _ in range(2000):
    trainer.step()
print(trainer.params["weight"].item())

3.1415926535897936

Both versions converge to the same value. Changing the backend of the new version is one keyword (device="mlx" or device="cuda", with float32 inputs).

15.4 Name Mapping

The most common replacements:

v0.1 name Replacement
trace_tensor(fn, specs) qb.jit(fn).lower(...) or qb.Compiler().trace(fn, specs)
tensor_jit_fn qb.jit(f)
tensor_jit_cuda_fn qb.jit(f, device="cuda")
tensor_jit_batch_fn qb.jit(f, max_traces=4)
tensor_value_and_grad_fn qb.value_and_grad(f, argnums)
tensor_value_and_grad_{mlx,cuda}_fn qb.jit(qb.value_and_grad(f, argnums), device="mlx" \| "cuda")
tensor_value_and_grad_batch_*_fn The same with max_traces=4
tensor_grad_scalar_fn qb.grad(f, argnums=(0, ..., n-1))
tensor_vjp_fn, tensor_jvp_fn qb.vjp(f, *primals), qb.jvp(f, primals, tangents)
tensor_jacobian_fn qb.jacobian(f, argnums=i)
tensor_hessian_scalar_fn qb.hessian(f, argnums=i)
tensor_hvp_scalar_fn qb.jvp(qb.grad(f), (x,), (v,))
tensor_vmap_fn qb.vmap(f, in_axes, out_axes)
tensor_vmap_{mlx,cuda}_fn qb.jit(qb.vmap(f, in_axes, out_axes), device=...)
tensor_vmap_{jvp,vjp}_*_fn qb.jvp/qb.vjp of qb.vmap(f), under jit(device=...)
tensor_vmap_hvp_scalar_fn qb.jvp(qb.grad(lambda x: qb.sum(qb.vmap(f)(x))), (x,), (v,))
tensor_cond qb.cond(pred, true_fun, false_fun, *operands)
tensor_fori_loop_region / tensor_fori_loop qb.fori_loop(..., operands=...) / with unroll=True
tensor_scan_region / tensor_scan qb.scan(f, init, length=n, operands=...) / with unroll=True
tensor_value_and_grad_data_parallel_cuda_fn qb.distributed.value_and_grad(f, devices=..., shard_argnums=...)
mlx_adam_loss_optimizer, cuda_adam_*_optimizer, cuda_adam_step qb.optim.Trainer(..., device="mlx" \| "cuda")
sum_gradients qb.tree.map(operator.add, a, b)
Adam qb.optim.Adam (init/update, or step)
Matrix, TraceGraph, grad_fn, jvp_fn, … qb.array and the core transforms (qb.grad, qb.jvp, qb.jit)

The complete mapping, with the status of each name, is Appendix A of the v0.2 API design.

15.5 Things That Behave Differently

This site is open source. Improve this page.