Numax v0.2.0 released with MAX 26.6

Community Showcase
1

numax: A lot of new features, closing the numpy + scipy APIs parity and improvements landed in v0.2.0, notably:

0. Every routine takes an owned Tensor or a borrowed View of one, so factoring a sub-block costs no copy.

from numax.core.tensorlike import View
from numax.linalg import cholesky, det, solve

var block = View(big.view().tile[2, 2](0, 0), big.context())
print(cholesky(block))         # factors the quadrant in place, no copy
print(det(block))

1. Full eigendecomposition, SVD, and Schur form now run on Tensor, device-resident, through MAX’s GEMM.

No more copying into a small Array to factor a matrix; it stays on the GPU the whole way.

from numax.linalg import eigh
from numax.prelude import *

def main() raises:
    var A = eye[f64, 4]() + ones[f64, 4, 4]()
    var result = eigh(A)                # blocked, device-resident
    print(result.values)                # ascending eigenvalues
    print(result.vectors)               # orthonormal columns

2. The same call differentiates, with no adjoint rule written anywhere.

eigh, cholesky, solve, qr are written once against a numeric trait; instantiate with Dual instead of Plain and the derivative falls out of the algorithm itself.

from numax.prelude import *

def ground_state[T: FloatLike](w: T) -> T:
    var H = build_hamiltonian(w)         # any FloatLike kernel
    return eigh[T, 24](H)[0][0]          # lowest eigenvalue

print(ground_state(Plain[f64].constant(1.0)))          # the value
print(ground_state(Dual[Plain[f64]].seed(1.0)).deriv)  # dE/dw, exactly

3. Batch thousands of small linear-algebra problems, one per GPU thread, from a single launch.

map_blocks hands each SIMD lane or GPU thread a whole factorization instead of one scalar.

from numax.core.tensor import map_blocks
from numax.linalg.array import cholesky_solve

def solve_step(problem: Array[Plain[f32], k_in]) -> Array[Plain[f32], k_out]:
    return cholesky_solve(problem)       # runs inside one lane

map_blocks[step=solve_step, gpu=True](device_a.view(), device_x.view(), gpu)
# 4096 independent 4x4 SPD solves, one launch

Check out the details as well as the examples.

2 Likes