v0.1.0

First release of νMAX, a numerical computing library built on MAX.

One kernel, several meanings

Every kernel in the library is written once, against the FloatLike trait.
The conformer you instantiate it with decides what the call returns:

Type	What a call gives you
Plain[dtype, width]	Ordinary SIMD at hardware speed
Dual[Inner]	f(x) and f'(x) from one call, forward-mode
Gradient[Inner, n_vars]	Every ∂f/∂xᵢ from one call
Compensated[dtype, width]	Roughly double dtype's precision, same algorithm
Decimal	Exact base-10 fixed point
Complex[Inner]	Complex results
Interval[Inner]	Rigorous enclosures

They nest. Gradient[Dual[Plain]] is a Hessian, Complex[Dual[Plain]]
differentiates holomorphically, Dual[Compensated] takes a derivative in
double-double. Nobody implemented autodiff for the incomplete beta function;
the chain rule lives in Dual's arithmetic rather than in any kernel, so
differentiating one costs nothing extra.

from numax import Dual, FloatLike, Plain, f32

def g[T: FloatLike](x: T) -> T:              # written once, no `dtype`
    return (-(x * x)).exp()

def main():
    print(g(Plain[f32](0.5)).v)              # 0.7788008
    var d = g(Dual[Plain[f32]].seed(0.5))
    print(d.value.v, d.deriv.v)              # 0.7788008 -0.7788008

One tensor, every device

Tensor owns a MAX DeviceBuffer, so the DeviceContext handed to a
factory decides host or device memory. Nothing else in the program changes,
and .view() yields the TileTensor every MAX kernel takes. The tensor
drivers (map, reduce, reduce_axis, broadcast_op_rows,
map_threaded) take one gpu: Bool compile-time parameter instead of
shipping two implementations, and map[gpu=True] launches through
DeviceContext.enqueue_function on Metal and CUDA with no per-backend
branch anywhere in the library.

The two-tier rule is what pays for that. Tier-1 kernels run a fixed
iteration count with no per-lane branching, doing selection arithmetically
instead of with an if, which is what makes them GPU-launchable unmodified.
Tier 2 is Plain-only and may loop to a tolerance. Every module and every
function declares which tier it is in, so a reader at a call site never has
to audit a body.

What is in it

Subpackage	Contents
numax.core	FloatLike and its conformers, Tensor, elementwise math, comparisons, sorting, the tensor engine
numax.special	Γ and B, erf, Bessel J/Y, Lambert W, elliptic K/E, orthogonal polynomials, activations
numax.linalg	cholesky, lu, qr, eigh, svd, solve, inverse, pinv, det, cond, norms, matmul
numax.optimize	newton/halley/bisection at fixed iterations; newton_tol, brentq, bfgs to a tolerance
numax.integrate	Gauss-Legendre/Simpson/trapezoid, rk4/dopri5; quad, quad_vec, solve_ivp
numax.interpolate	Horner, cubic splines, Chebyshev fits
numax.fft	fft/ifft, rfft/irfft, fft2, fftfreq, fftshift
numax.signal	convolve, correlate, Hann/Hamming/Blackman windows
numax.stats	Descriptive statistics, nine scipy.stats-shaped distribution namespaces, sampling
numax.io	.npy read/write, byte-identical to numpy.save, plus numax's own NMX1 format

Subpackages carry NumPy and SciPy names, so an import has an obvious
counterpart. from numax.prelude import * is the common surface in one
line. It leaves out every name that would shadow a Mojo builtin, because a
module-level definition replaces the builtin rather than overloading it.

The .npy support is small but it removes a conversion step: a program
ported from NumPy can read the files it already has and hand results back
the same way.

Performance

Every number below was measured, and CPU and GPU results are never combined
into one comparison. On an A10G's GPU, numax reaches 60,955 M elem/s against
torch.compile's 53,607, about 83% of the card's bandwidth spec. On the same
machine's CPU, map over a Plain kernel runs at 0.998x a hand-written
raw-SIMD loop and 4.5x NumPy. Every row in docs/performance.md compares
implementations on one processor.

Accuracy

Every approximation documents an error bound, and pixi run accuracy checks
each one against mpmath references computed at 50 digits and checked into
the repository.

Limits worth knowing before you start

Shapes are compile-time parameters. There is no broadcasting, no strided
views, no slicing and no fancy indexing. Those wait on a runtime-shape
owner, because a view over a slice is not row-major and every driver in
numax.core.tensor requires that it is.

Above the tensor drivers the surface is not uniformly rank-generic yet.
numax.stats reduces over every element with no axis=; sorting flattens;
transpose is 2-D; concatenate, split and stack are rank-1. The
SciPy-shaped algorithms are fixed-size Array kernels, so linalg works on
matrices and fft2 on a square transform, while quad, solve_ivp, the
splines and the distributions are 1-D.

Anything whose length depends on a value rather than a parameter has to say
so in its signature. unique returns a full-length tensor plus a count
rather than a right-sized array, and its docstring tells you to read only
the first count entries.

.npz is out of scope. It is a zip container, and while stored-only entries
would be easy, deflate is not.

Running the GPU examples needs a real device. CI compiles them against a
pinned target so their call sites stay honest, but execution is a local
check, verified by hand on both Metal and CUDA.

This is a young library and the APIs may change.

Install

[workspace]
channels = ["https://conda.modular.com/max", "conda-forge"]
preview = ["pixi-build"]

[dependencies]
numax = { git = "https://github.com/ehsanmok/numax.git", tag = "v0.1.0" }

Requires pixi. mojo ==1.0.0 and max ==26.5 come in
transitively. numax is distributed as source rather than a .mojopkg,
because a package embeds the compiler version it was built with and refuses
to load against any other.

Verified before release on linux-64: the git dependency resolves from this
tag, builds, and runs the example above with no import path flag. The
manifest declares linux-64, linux-aarch64 and osx-arm64; CI runs on linux-64
and osx-arm64.

Verification

44 test suites and 664 tests, 14 CPU examples, the accuracy harness against
the mpmath references, the CPU benchmark, mojo doc over the whole public
generic surface, and a formatting check. All green on the tagged commit, on
both ubuntu-latest and macos-15.

Full documentation: features ·
architecture · parity ·
performance