Rune vs. JAX — A Practical Comparison

This guide explains how rune's transformations relate to JAX, focusing on:

• How core concepts map (grad, vjp, jvp, vmap, pytrees, custom rules, remat)
• Where the APIs feel similar vs. deliberately different
• What rune does not have yet, stated honestly

If you already use JAX, this should be enough to become productive in rune quickly.

1. Big-Picture Differences

2. Pytrees → Ptree.S

This is the deepest difference. JAX flattens arbitrary registered containers into lists of leaves at runtime:

import jax

params = {"w": w, "b": b}          # any registered pytree
grads = jax.grad(loss)(params)      # same pytree of gradients

Rune has no runtime tree. A parameter structure is a record you define, and you tell the library how to traverse its tensor leaves by implementing Nx.Ptree.S — three hand-written one-liners, no ppx, no registration table:

type params = { w : Nx.float32_t; b : Nx.float32_t }

module Params = struct
  type t = params

  let map (f : 'a 'b. ('a, 'b) Nx.t -> ('a, 'b) Nx.t) { w; b } =
    { w = f w; b = f b }

  let map2 (f : 'a 'b. ('a, 'b) Nx.t -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t) p q =
    { w = f p.w q.w; b = f p.b q.b }

  let iter (f : 'a 'b. ('a, 'b) Nx.t -> unit) { w; b } =
    f w;
    f b
end

Every transformation takes the module as a first-class argument and preserves the type: the gradient of a function of params is a params. There is no tree_map because Params.map is tree_map, specialized to your type — and Rune.Ptree (that is, Nx.Ptree.t) is the stock dynamic instance for structures only known at runtime, the closest analogue of a raw pytree.

Where JAX distinguishes leaves by position in a flattened list, rune leaves keep their record field names, dtypes, and shapes in the type. Mixed dtypes work: a single backward pass produces a gradient for every leaf, each with its leaf's dtype.

3. Reverse Mode

JAX

def loss(params):
    return jnp.mean((x @ params["w"] + params["b"] - y) ** 2)

grads = jax.grad(loss)(params)
loss_value, grads = jax.value_and_grad(loss)(params)

rune

let () =
  let x = Nx.create Nx.float32 [| 4; 3 |] (Array.init 12 float_of_int) in
  let y = Nx.create Nx.float32 [| 4; 1 |] [| 0.; 1.; 2.; 3. |] in
  let loss p =
    Nx.mean (Nx.square (Nx.sub (Nx.add (Nx.matmul x p.w) p.b) y))
  in
  let params =
    { w = Nx.zeros Nx.float32 [| 3; 1 |]; b = Nx.zeros Nx.float32 [| 1 |] }
  in
  let grads = Rune.grad (module Params) loss params in
  let loss_value, grads' =
    Rune.value_and_grad (module Params) loss params
  in
  ignore (grads, loss_value, grads')

Both require a scalar output. JAX's argnums has no equivalent: differentiate with respect to the parameter structure and close over everything else. For a function of one tensor, grad' skips the module argument.

Auxiliary outputs

JAX uses a flag; rune has dedicated _aux variants:

(loss, aux), grads = jax.value_and_grad(f, has_aux=True)(params)

let loss, grads, aux = Rune.value_and_grad_aux (module Params) f params

4. VJP and JVP

JAX's jax.vjp returns a pullback closure; rune offers both that shape (vjp_fun) and a one-shot version (vjp) that takes the cotangent directly:

y, pullback = jax.vjp(f, x)
grads = pullback(ct)

let () =
  let f v = Nx.mul v v in
  let x = Nx.create Nx.float32 [| 3 |] [| 1.; 2.; 3. |] in
  let ct = Nx.ones Nx.float32 [| 3 |] in

  (* One-shot: *)
  let y, g = Rune.vjp' f x ct in

  (* Reusable pullback, as in JAX: *)
  let y', pullback = Rune.vjp_fun' f x in
  let g' = pullback ct in
  ignore (y, g, y', g')

Forward mode is nearly identical in both:

y, tangent = jax.jvp(f, (x,), (v,))

let () =
  let f v = Nx.mul v v in
  let x = Nx.create Nx.float32 [| 3 |] [| 1.; 2.; 3. |] in
  let v = Nx.ones Nx.float32 [| 3 |] in
  let y, tangent = Rune.jvp' f x v in
  ignore (y, tangent)

For functions returning a structure rather than one tensor, use vjp2/jvp2 with a second module describing the output.

5. vmap and Per-Sample Gradients

f_batched = jax.vmap(f)                      # map axis 0 of every input
jax.vmap(f, in_axes=(0, None))               # hold the second input fixed
per_sample = jax.vmap(jax.grad(loss))(batch) # per-example gradients

in_axes translates directly: Some i for a mapped axis, None for a constant, one entry per leaf in traversal order:

let () =
  let f v = Nx.sum (Nx.mul v v) in
  let batch = Nx.ones Nx.float32 [| 10; 5 |] in
  let results = Rune.vmap' f batch in
  ignore results (* shape [10] *)

(* Hold the second leaf fixed: *)
let ys =
  Rune.vmap ~in_axes:[ Some 0; None ] (module Pair) f pairs

(* Per-sample gradients: vmap2 of grad. *)
let per_sample =
  Rune.vmap2
    (module Example)
    (module Params)
    (fun ex -> Rune.grad (module Params) (loss ex) params)
    batch

Two honest caveats relative to jax.vmap:

• rune's vmap has no batching rule for the fft family and matrix decompositions (cholesky, qr, svd, eig, eigh); those raise on batched inputs.
• Implicit RNG inside the mapped function draws identical values for every lane (JAX avoids this by making you thread keys; in rune, thread randomness in as mapped inputs).

6. Custom Rules

The correspondence is direct. JAX:

@jax.custom_vjp
def f(x): return jnp.square(x)

def f_fwd(x): return jnp.square(x), x
def f_bwd(res, ct): return (ct * 2 * res,)
f.defvjp(f_fwd, f_bwd)

rune packs the same three pieces into one call — the forward function returns the residual alongside its result:

module Vec = struct
  type t = Nx.float32_t

  let map (f : 'a 'b. ('a, 'b) Nx.t -> ('a, 'b) Nx.t) v = f v
  let map2 (f : 'a 'b. ('a, 'b) Nx.t -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t) = f
  let iter (f : 'a 'b. ('a, 'b) Nx.t -> unit) v = f v
end

let f x =
  Rune.custom_vjp
    (module Vec)
    ~fwd:(fun x -> (Nx.square x, x))
    ~bwd:(fun res ct -> Nx.mul ct (Nx.mul_s res 2.0))
    x

let () =
  let x = Nx.create Nx.float32 [| 3 |] [| 1.; 2.; 3. |] in
  Nx.print_data (Rune.grad' (fun v -> Nx.sum (f v)) x)

custom_jvp mirrors jax.custom_jvp the same way. One difference to know: in rune a custom_vjp raises if differentiated in forward mode (and vice versa) — define both rules if you need both modes, where JAX can sometimes transpose a JVP rule automatically.

7. Gradient Checkpointing

jax.checkpoint (a.k.a. jax.remat) maps to remat:

y = jax.checkpoint(expensive)(x)

let () =
  let expensive v = Nx.mean (Nx.square (Nx.sin v)) in
  let x = Nx.create Nx.float32 [| 3 |] [| 1.; 2.; 3. |] in
  let g =
    Rune.grad' (fun v -> Rune.remat (module Vec) expensive v) x
  in
  ignore g

Both recompute the wrapped function during the backward pass instead of retaining its intermediates; gradients are unchanged.

8. Control Flow

In JAX, Python control flow breaks under jit, so lax.cond/lax.scan/lax.while_loop are mandatory inside compiled functions. In rune everything runs eagerly, so ordinary OCaml control flow works inside every transformation:

let () =
  let f x = if Nx.item [] x > 0.0 then x else Nx.neg x in
  ignore (Rune.grad' f (Nx.scalar Nx.float32 2.0))

Rune still provides scan, cond, and while_loop — not because you need them today, but because their signatures are staging-ready: code written with them differentiates and vectorizes now, and a future jit can stage them as structured control flow instead of unrolled traces. lax.scan's carry-and-stacked-outputs contract translates directly:

final, ys = jax.lax.scan(f, init, xs)

let final, ys = Rune.scan (module Carry) ~f ~init xs

9. Jacobians, Hessians, HVPs

JAX's docs derive the Hessian-vector product as jvp of grad; rune ships that composition as hvp, matrix-free, for any parameter structure.

10. Gradient Checking

from jax.test_util import check_grads
check_grads(f, (x,), order=1)

let () =
  let f v = Nx.sum (Nx.mul v v) in
  let module V64 = struct
    type t = Nx.float64_t

    let map (f : 'a 'b. ('a, 'b) Nx.t -> ('a, 'b) Nx.t) v = f v
    let map2 (f : 'a 'b. ('a, 'b) Nx.t -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t) = f
    let iter (f : 'a 'b. ('a, 'b) Nx.t -> unit) v = f v
  end in
  match
    Rune.check_grads
      (module V64)
      f
      (Nx.create Nx.float64 [| 3 |] [| 1.; 2.; 3. |])
  with
  | Ok () -> print_endline "ok"
  | Error msg -> print_endline msg

Both compare autodiff against finite differences along directions rather than element by element. Use float64 for reliable results.

11. Randomness

JAX threads explicit splittable keys. Rune uses Nx's implicit scoped RNG: wrap the program in Nx.Rng.run ~seed for reproducibility, and Nx.rand/Nx.randn draw from the ambient scope:

let () =
  Nx.Rng.run ~seed:0 @@ fun () ->
  ignore (Nx.randn Nx.float32 [| 3 |])

The trade-off surfaces under vmap: with explicit keys you would pass one key per lane; with the implicit scope, in-function draws are identical across lanes, so per-lane randomness must be a mapped input.

12. What Rune Does Not Have (Yet)

Rune's failure model is deliberate: operations without a rule raise Invalid_argument rather than silently producing zero or wrong gradients.

13. Quick Cheat Sheet