nx

NumPy's power. OCaml's type safety. Zero compromises.

why nx?

your numpy code, but safer

The API feels like NumPy because it is like NumPy. But the compiler catches your dimension mismatches.

actually zero-copy

Slicing returns views. Broadcasting doesn't allocate. Reshaping is free. We mean it.

backends that matter

CPU today. Metal on macOS. CUDA coming. Same code runs everywhere.

part of something bigger

First-class citizen in Raven. Your arrays flow seamlessly into plotting, autodiff, and neural networks.

show me the code

NUMPY

# Create and reshape
a = np.arange(12).reshape(3, 4)

# Slice (might copy)
b = a[:, 1:3]

# Reduction
c = a.mean(axis=0)

# This crashes at runtime
x = np.ones((2, 3))
y = np.ones((3, 2))
z = x + y  # ValueError!

NX

(* Create and reshape *)
let a = arange 12 |> reshape [|3; 4|]

(* Slice (never copies) *)
let b = get_slice [All; R [1; 3]] a

(* Reduction *)
let c = mean ~axes:[|0|] a

(* This would fail at compile time *)
let x = ones float32 [|2; 3|]
let y = ones float32 [|3; 2|]
(* let z = add x y  -- won't compile! *)

the good parts

Types that work
float32, float64, int32, int64, uint8, and more. Generic functions that preserve types.

Broadcasting that makes sense
[|3; 1; 5|] + [|1; 4; 5|] = [|3; 4; 5|]. Same rules as NumPy.

I/O that just works
Load your .npy files. Save as images. Interop with Python when you need to.

get started

Nx isn't released yet. For now, check out the documentation to learn more.

When it's ready:

# Install
opam install nx

# Try it
open Nx

let () = 
  let x = rand float32 [|3; 3|] in
  print_tensor x