Coming from Python

Library mapping

Key differences

Explicit types

NumPy casts types silently. Nx does not.

# Python: silently upcasts int + float -> float
a = np.array([1, 2, 3])
b = a + 1.5  # works

(* OCaml: types must match *)
let a = Nx.create Nx.int32 [| 3 |] [| 1l; 2l; 3l |]
(* Nx.add a (Nx.scalar Nx.float32 1.5)  -- type error *)

(* Cast explicitly *)
let a_f = Nx.cast Nx.float32 a
let b = Nx.add_s a_f 1.5

Array literals

NumPy uses Python lists. Nx uses OCaml arrays with [| |] syntax.

x = np.array([[1, 2], [3, 4]])

let x = Nx.create Nx.float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |]

Slicing

NumPy uses [] with :. Nx uses the slice function with index constructors, which are Nx's, so a slice is written inside a local open.

x[0:2, :]           # first two rows
x[:, 1]             # second column
x[::2]              # every other row

Nx.(slice [ R (0, 2); A ] x)          (* first two rows *)
Nx.(slice [ A; I 1 ] x)               (* second column *)
Nx.(slice [ Rs (0, dim 0 x, 2) ] x)   (* every other row *)

No separate tensor type

In PyTorch, torch.Tensor is different from numpy.ndarray. In Raven, Rune operates directly on Nx.t values. There is no wrapper type.

# PyTorch: convert between types
x_np = np.array([1.0, 2.0])
x_torch = torch.from_numpy(x_np)
x_torch.requires_grad_(True)

(* Raven: just use Nx tensors directly *)
let x = Nx.create Nx.float32 [| 2 |] [| 1.0; 2.0 |]
let gradient = Rune.grad' (fun x -> Nx.sum (Nx.mul x x)) x

Functional transformations

JAX users will find Rune familiar. PyTorch users: think of grad as a function transformer, not a method on tensors.

# JAX style
grad_fn = jax.grad(loss_fn)
grads = grad_fn(params)

# PyTorch style
loss = loss_fn(params)
loss.backward()
grads = params.grad

(* Rune: JAX-style functional transforms over your own typed record *)
let params_ptree = Nx.Ptree.instantiate (module Params)
let grads = Rune.grad params_ptree loss_fn params

(* Or compute value and gradient together *)
let loss, grads = Rune.value_and_grad params_ptree loss_fn params

Where JAX registers pytrees, Rune takes a structure: Params implements Nx.Ptree.S with one function, walk, over your record's fields, Nx.Ptree.instantiate makes it a value, and gradients come back with the same type as the parameters.

Module-based layers

Kaun layers are plain records with init and apply functions, not classes with forward — and a model is a record you define, not a container object.

# PyTorch
class Model(nn.Module):
    def __init__(self):
        self.linear = nn.Linear(784, 10)
    def forward(self, x):
        return self.linear(x)
model = Model()

(* Kaun: a model is a typed record of layer records *)
type ('a, 'd) model = { linear : ('a, 'd) Kaun.Linear.t }

let apply p x = Kaun.Linear.apply p.linear x

let params =
  { linear = Kaun.Linear.init ~key:(Nx.Rng.key 0) ~inputs:784 ~outputs:10 }

Parameters are plain data (records of Nx tensors), not hidden inside objects. A constructor draws them from the key it is given.

Random numbers

NumPy and PyTorch draw from a global generator that a seed resets. Nx draws from keys: a draw is a function of its key, dtype and shape. A sampler takes ~key, or else the next key of the enclosing Nx.Rng.with_key scope; outside every scope, the next key of the domain's own, seeded from the system's entropy. Nx.Rng.split makes independent keys from one (nx.mli, Rng).

np.random.seed(0)
x = np.random.normal(size=3)

let x = Nx.Rng.normal ~key:(Nx.Rng.key 0) Nx.float32 [| 3 |]

let y =
  Nx.Rng.with_key (Nx.Rng.key 1) (fun () -> Nx.Rng.uniform Nx.float32 [| 3 |])

Devices

In PyTorch, a tensor's device is a runtime attribute. In Raven, a value's type carries its device set. A value made from numbers is a host value, of type (_, _, Nx.host) Nx.t, and Nx.devices mints a set with a brand of its own from devices that rig opens. Nx.place moves a value onto a set, as Nx.Ptree.place moves a model. A set minted without ~kernels over GPUs computes nothing eagerly, and an operation there raises, so a function over GPU values runs compiled, under Rune.jit, one device per compiled program (nx.mli, Device sets; rune.mli, jit).

x = torch.tensor([1.0, 2.0]).to("cuda")
y = x * 2

module Gpu = (val Nx.devices [ Result.get_ok (Rig.open_ (Rig_cuda.gpu 0)) ])

let x = Nx.place Gpu.on (Nx.create Nx.float32 [| 2 |] [| 1.0; 2.0 |])
let double = Rune.jit' (fun x -> Nx.mul_s x 2.)
let y = double x (* a Gpu.d value *)

PyTorch updates a tensor in place, as x.add_(1) does. Nx never changes a value. A program that wants a result to reuse a value's memory gives the value up at the call. Nx.donate x, or Nx.Ptree.donate for a structure, lets the operation or compiled call that reads it write its result there, and the donated value is dead afterwards (nx.mli, Donation).

DataFrames

pandas computes each operation as it is called, and finds a wrong column name when the line that uses it runs. Talon builds a query from typed column handles, checks every verb against its input's columns before reading any data, and computes the query when it runs.

# pandas
df["bmi"] = df["weight"] / df["height"] ** 2

(* Talon: a typed expression in a query *)
let weight = Talon.Col.float "weight" and height = Talon.Col.float "height"

let with_bmi df =
  Talon.Query.(df |> derive Talon.Expr.[ "bmi" := weight /. (height *. height) ])

Detailed comparisons

Each library has a dedicated comparison page with side-by-side code examples:

• Nx vs NumPy
• Kaun vs PyTorch/Flax
• Vega vs Optax
• Brot vs HuggingFace Tokenizers
• Talon vs pandas
• Sowilo vs OpenCV (contrib)
• Fehu vs Gymnasium (contrib)

edit this page ↗