numax v0.3.0

0.3.0 makes every Tensor-tier routine run on the accelerator, adds four SciPy areas numax never had, and renames a lot of the surface so one word means one thing. It is a hard break: the renames below land without deprecation aliases. 173 commits since v0.2.0, 1,824 tests in 157 suites, and the whole gate (tests, examples, accuracy, type-check, compile-fail checks, format, CI) ran on Linux with an NVIDIA A10G as well as on the Apple M3 Pro the release was developed on.

Breaking changes

If you are upgrading from 0.2.0, these are the edits you will make.

Old	New
View, ViewOver	TensorView, TensorViewOver (now exported)
Tensor.view()	Tensor.tile()
static_view, dynamic	as_static, as_dynamic
numax.core.array	numax.core.tensor
numax.core.tensor (the map/reduce drivers)	numax.core.functional
TensorLU, TensorQR, TensorSVD, TensorEigh, TensorSchur, TensorMinimizeResult, TensorIVPResult and the other Tensor-prefixed result structs	LU, QR, SVD, Eigh, Schur, MinimizeResult, IVPResult, ...
scalar-tier IVPResult, register-tier MinimizeResult	ScalarIVPResult, ArrayMinimizeResult
Chebyshev, CubicSpline in the register tier	ArrayChebyshev, ArrayCubicSpline
norm selectors inf, fro, neg_inf	NORM_INF, NORM_FRO, NORM_NEG_INF
signal modes full, same, valid	MODE_FULL, MODE_SAME, MODE_VALID
count-based arange[num](start, step)	arange_n; arange now takes NumPy's stop or start, stop, step

Behavior changes that are not renames:

• Tensor's constructors take the DeviceContext last and optional, like the factories. A wrong-length initializer raises "Tensor: N values for M elements" instead of reading past the end.
• Tensor.buffer and layout, and TensorView's tile and device, are private. Reach the data through to_host and copy_from_host.
• solve_triangular's upper now names the triangle that is stored.
• The six register-tier subpackages (linalg, fft, signal, integrate, optimize, interpolate) no longer have a public .array module. See "One name per operation".
• sort, ravel and diagflat on a run-time-shaped tensor no longer ask MAX for a buffer of -1 elements; a run-time-shaped input gives a run-time-shaped result.

Everything runs on the accelerator

Every public Tensor-tier routine now has a device path. Where MAX ships a kernel numax calls it. Where MAX does not, numax writes one in MAX's idiom and records the gap in docs/parity.md. Host-only is no longer something a routine may ship with.

What changed on the device:

• The spectral chain (sytrd, eigvalsh, eigh, hessenberg, eigvals, schur, gebrd, svdvals, svd, matrix_rank, and everything built on them: the matrix functions, pinv, orth, null_space, polar, cond, lstsq(method="svd")) works at gpu=True. In 0.2 that was a compile error, because sytrd[gpu=True] disagreed with the host at n = 4, block = 1. That is fixed and each routine has a device test.
• The operators (a + b and the rest) follow the tensor to wherever it lives. So do the shape routines (reshape, concatenate, stack, split, tile, repeat, roll, flip, broadcast_to, diag, rot90, and so on) and the creation functions (zeros, ones, full, eye, linspace, arange).
• Scans: cumsum, cumprod, cumulative_trapezoid. MAX's nn.cumsum is CPU-only, so this is numax's own blocked scan.
• Sorting and selection: sort, argsort, median, quantile, percentile, iqr, searchsorted, unique, rankdata. The sort is numax's own bitonic sort, because MAX's GPU argsort returns wrong results past 256 elements on Metal. Axis-wise median and mode sort every slice at once.
• Statistics: histograms (histogram, histogram2d, histogramdd, bincount, digitize) by atomics into per-block bins, the descriptive statistics, pearsonr, linregress, spearmanr, kendalltau, and every hypothesis test except ks_1samp, whose Float64 cdf parameter Metal cannot compile.
• Data-dependent output length: nonzero, argwhere, extract, compress, find_peaks, peak_prominences, by mask, scan and scatter. The one scalar read back sizes the result.
• Recurrences: lfilter, filtfilt, sosfilt, decimate, resample, hilbert, istft, savgol_filter, and spline construction.
• Banded solvers: solve_banded and solveh_banded for tridiagonal systems, by parallel cyclic reduction.
• Iterative drivers: least_squares and curve_fit keep residuals and the Jacobian on the device, and solve_ivp no longer stages state through the host.
• The Array tier's tier-2 routines (the optimize drivers, pinv, firwin) no longer need float64 or raises in their bodies, so they run inside a kernel.

The exceptions are listed per routine in docs/parity.md: a single banded Cholesky or Toeplitz system, the narrow-band iteration inside the eigen and SVD routines (its O(n^2) accumulation is on the device), and ks_1samp.

How it is checked: tests/ still never names gpu=True, because a GPU-less CI runner cannot compile the device path. A separate tests_gpu/ tree compares every gpu=True routine against its host answer at float32. It passed on Metal and, in this release, on an A10G (all 16 areas). CI compiles it for sm_80 on every push.

Four new SciPy areas

numax.spatial has distance.cdist, pdist and squareform; Euclidean goes through the Gram-matrix identity on linalg.matmul, the other metrics are one elementwise pass each. KDTree with query and query_ball_point, built as a linear BVH over Morton-ordered points by a device sort and queried one point per thread. Delaunay, ConvexHull, Voronoi and Rotation are not here.

numax.ndimage has convolve, correlate, gaussian_filter, uniform_filter, median_filter, map_coordinates, shift, zoom, label, binary_erosion, binary_dilation and distance_transform_edt. It is imported as numax.ndimage, not from the root, because its convolve and correlate are numax.signal's names with a different meaning.

numax.cluster has vq.vq, kmeans and kmeans2 as Lloyd's iteration over cdist and MAX's ArgMin, and hierarchy.linkage and fcluster. linkage runs the primitive algorithm with the distance matrix and each merge's row update on the device; it is not the nearest-neighbor-chain version.

numax.optimize gains linprog (SciPy's former homogeneous self-dual interior point, with the normal matrix formed and factored on the device), milp (branch and bound over it), and minimize(method="slsqp") with LinearConstraint and NonlinearConstraint. trust-constr and the global optimizers are not here.

Examples are clustering.mojo (spatial and cluster), image_processing.mojo and constrained_optimization.mojo under examples/intermediate/.

NumPy and SciPy gaps closed

• Core: unique_counts and unique_inverse (NumPy 2's split of unique), isin, intersect1d, union1d, setdiff1d, nan_to_num, logaddexp, logaddexp2, sinc, heaviside, gcd, lcm, the bitwise operations and shifts, nanargmax, nanargmin, ediff1d, unwrap, and keepdims on the axis reductions.
• fft: fftn, ifftn, rfftn, irfftn, irfft2, hfft, ihfft, dctn, idctn, dstn, idstn, and norm= on the transforms. The engine needed no axis argument: each pass transforms the last axis and writes through a rotated view. rfftn and irfftn stop at rank 3.
• signal: Second-order sections (output=OUTPUT_SOS on every design, plus tf2sos, zpk2sos, tf2zpk, sos2tf, sosfreqz, sosfilt_zi, sosfiltfilt), the order estimators buttord, cheb1ord, cheb2ord, ellipord, iirnotch, iirpeak, bilinear, group_delay, lfiltic and lfilter(zi), upfirdn, resample_poly, five more windows (tukey, gaussian, flattop, nuttall, chebwin), chirp, sawtooth, square, peak_widths, argrelextrema, convolve2d and correlate2d.
• stats: .rvs() on every distribution by inverse CDF over a shared Generator; Generator.choice, permutation, shuffle, gamma, beta, poisson, binomial, lognormal, multivariate_normal and spawn; mean, var, std, interval and entropy on every distribution; new distributions uniform_dist, lognorm, weibull_min, cauchy, laplace, rayleigh, logistic, pareto, bernoulli, geom, nbinom, hypergeom and multivariate_normal; tests kruskal, levene, bartlett, fisher_exact, binomtest, chi2_contingency, skewtest, kurtosistest, normaltest, jarque_bera, bootstrap and permutation_test; and median_abs_deviation, ecdf and gaussian_kde.
• special: expit, log_expit, ndtr, log_ndtr, ndtri, gammaincinv, betaincinv, polygamma, ellipkinc, ellipeinc, elliprf, elliprd and wofz. elliptic_k and elliptic_e are now computed through Carlson's forms. betainc's lower tail reaches below 1e-30.
• linalg: ldl, generalized eigh(a, b) and eigvalsh(a, b), solve_sylvester, solve_continuous_lyapunov, solve_discrete_lyapunov, solve_continuous_are, and lu returning P, L, U.
• integrate: solve_bvp, cumulative_simpson, romb, and t_eval, dense output and events on solve_ivp.
• interpolate: General-degree BSpline and make_interp_spline, interpn in up to eight dimensions, RBFInterpolator, NearestNDInterpolator and griddata (nearest-neighbor only, since there is no Delaunay).

Masked arrays, datetime64 and structured dtypes, the numpy.polynomial classes, and sparse and Krylov solvers are recorded as absent, each with its reason, in docs/parity.md.

Three accuracy fixes came with this: log1p, log2, exp2 and cosh at float64 now compute to double precision, max and min of all-infinite input return the infinity (MAX's ReduceMax identity was min_finite), and the _hqr Francis iteration adopts dlahqr's shift schedule so it cannot cycle at float32.

A surface that says what it means

• Operators: a + b works across any TensorLike operands, so Static + Dynamic, Tensor + TensorView and a broadcast + compile. There are reflected forms (2.0 * a), in-place forms (a += b), a @ b, a ** b, and the comparisons < <= > >= == !=, which return a bool tensor. TensorView has the same operators and the same printing as Tensor.
• Indexing: a[i:j] and a[i:j, k:l] return a TensorView with no copy, a.block[m, n](r, c) returns a view of a compile-time shape, rank-3 scalar indexing works, and every scalar index is bounds-checked. slice() stays as the copying spelling.
• Errors: every raise starts with the routine's name. An operator dtype or shape mismatch, an @ with mismatched inner extents, and an Array-tier matrix too large to unroll each stop on a message that names the fix instead of "no matching overload" or a compile that never ends. tests_compile_fail/ holds each message and pixi run compile-fail-check runs in CI.
• The device fallback has a policy: set_fallback("warn" | "raise" | "silent"); "warn" is the default, and "raise" lets a test or an agent turn a residency mismatch into an error.
• Names: one word, one meaning: view is TensorView and nothing else, tile is MAX's TileTensor, array is the register tier. a.tile() is the only way to a TileTensor, and no basic example spells it.

One name per operation

0.2 had a Tensor cholesky and an Array cholesky under different imports. Now cholesky, solve, fft, quad and the rest are one function each, overloaded on the argument: a Tensor takes the tier that goes through MAX, an Array[T, n] takes the differentiable, register-resident one. The register bodies moved to private _array/ packages. A caller of one tier pays no compile time for the other, and mojo doc shows both. Both bodies survive on purpose, because they are different algorithms (blocked and device-resident against an unblocked register loop).

Matrices whose size is read at run time

The dense routines no longer need n in the type. lu_factor, solve, det and slogdet (over DynamicLU), cholesky, qr_factor and lstsq (over DynamicQR), matvec, eigvalsh, eigh, svdvals, svd, eigvals and schur (results DynamicEigh, DynamicSVD, DynamicEigenvalues, DynamicSchur) all take a Dynamic tensor. Each static overload wraps a run-time-shaped body, and the two agree to the bit. A matrix loaded with numpy.load_dyn can now be factored. hessenberg, sytrd, gebrd and what builds on the spectral routines (pinv, the matrix functions, the matrix-equation solvers) are still static.

Performance

Every number is one processor, float32, and CPU and GPU are never in the same comparison. docs/performance.md has the harness for each, the M3 Pro sections from development, and a new Linux x86_64 (AMD EPYC 7R32, 64 vCPU) and NVIDIA A10G section measured for this release.

On the NVIDIA A10G, the Gaussian sweep at 2^26 elements runs at 60,939 M elem/s per call (61,519 amortized), against a hand-written CuPy kernel at 60,275 (61,037), torch.compile at 53,734 (56,199) and PyTorch eager at 19,801 (19,921). The copy kernel reaches 489 GB/s and the Gaussian 489, so the exp is free and the kernel sits at 81% of the card's spec. MAX's GEMM reaches 21,848 GFLOP/s at n = 1024 against cuBLAS's FP32 path at 14,924 through PyTorch. The factorizations around it are far behind cuSOLVER: at n = 1024, cholesky 50.6 against 598 GFLOP/s, lu_factor 44.6 against 282, qr_factor 17.3 against 269. The panels are the ceiling, not the GEMM.

On the x86 CPU (EPYC 7R32), map runs the Gaussian at 1,418 M elem/s and map_threaded at 8,298, against NumPy at 314, thermite at 1,334 and torch.compile at 1,813. lu_factor at n = 4096 is 79.3 GFLOP/s against SciPy's 49.5, because OpenBLAS's getrf falls off there. Elsewhere LAPACK is ahead: at n = 1024, cholesky 17.3 against 52.3, lu_factor 17.0 against 46.1, qr_factor 4.0 against 15.5, and the spectral routines reach 0.17 to 0.43 of SciPy. numax is well ahead on interp (3.07 ms against 48.2), norm.cdf at 2^24 (10.0 ms against 485.8), cov (6.0 against 26.7), welch (33.7 against 249.4) and medfilt (1.02 against 9.34).

On the Apple M3 Pro, the spectral routines at n = 1024 reach 0.41 (eigvalsh), 0.45 (eigh), 0.34 (svdvals), 0.70 (eigvals) and 0.45 (schur) of Accelerate's LAPACK. svd is at 0.056 and is the one target this release missed. svdvals and eigvalsh reduce in two stages through a band, eigvals and schur run LAPACK's multishift QR above n = 75, and the panel kernels fold their cross-thread reductions through threadgroup trees.

Where numax loses: on x86, fft at 2^20 is 0.34 of pocketfft, rfft is 0.12, fft2 is 0.08, fftconvolve is 8x slower and quantile is 0.32. These are measured gaps with no fix in 0.3. Device dot and nrm2 trail PyTorch.

One regression was caught before it shipped: making the blocked factorizations' wide loads safe on CUDA by forcing element alignment cost the A10G Cholesky 20%. The shipped fix checks the address at run time and takes the wide path whenever it is aligned.

What the A10G found

0.3 was developed on macOS. Running the gate on an A10G found three defects Metal hides, each fixed with a tests_gpu/ test and logged upstream with a standalone repro under .cursor/upstream-repros/:

• TileTensor's wide loads and stores assume vector alignment on NVIDIA, so a sub-block at an odd corner or stride faulted (pack_block, the trsm dot product, every trailing-update epilogue).
• firwin's center tap was sin(eps) / eps, and NVIDIA's float32 sin is sin.approx.ftz.f32, which returns 0 there. The tap is now taken at its defining value.
• lgamma's reflection side took ln|sin(pi x)| at integers, which is -inf for the same reason, and the discarded lane carried 0 * -inf into the blend. nbinom, hypergeom and t.rvs returned NaN at integer arguments.

Two things about the CUDA path are different from Metal and worth knowing. float32 sin and cos are the hardware approximations, so quantum_well.mojo's GPU and CPU results differ by 1.4e-5 on CUDA where they are bit-identical on Metal. And float64 linalg.matmul does not compile for a device at the 26.6 pin, so every routine that goes through it (matmul, cholesky, solve, eigvalsh, svdvals and the rest) is float32-only on the device. float64 device work is limited to elementwise kernels.

Documentation

llms.txt at the repo root is the agent guide, docs/quickstart.md goes from install to a printed answer in under 30 lines, and public routines in core and every prelude name carry Args, Parameters, Returns and Raises sections. Every module and public function declares its tier. pixi run doc-check compiles every mojo block in llms.txt, the quickstart and the README, and pixi run typecheck (mojo doc over the whole generic surface) runs in CI. Blocks that open a DeviceContext or ask for gpu=True compile for sm_80, so a GPU-less runner can check them. The README, the root docstring and recipe.yaml say the same thing again against the 0.3 surface.

Known limits

• svd and eigh still pay the one-stage reduction. svd is 0.056 of LAPACK on the M3 Pro and 0.17 on the EPYC.
• svdvals at float32 is accurate to 7.9e-7 against Accelerate's 3.8e-8.
• Metal cholesky is 0.26 to 0.52 of MPS; what is left is per-launch overhead.
• float64 on Metal does not exist (no double), and float64 device matmul does not compile on CUDA at this pin.
• ROCm is unmeasured.
• The FFT, fftconvolve, convolve and quantile gaps on x86 above.
• Tensor is dtype-monomorphic. Derivatives, bounds and extra precision through an algorithm are the Array tier's.

Install

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

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

Requires max-core ==26.6 and mojo >=1.0.0,<2.0.0, both pulled in by pixi. This install was resolved in a fresh project from the tag before this release was published.