mojolearn.linalg

Float32 linear algebra: a bit-identical matrix product, low-bit products and NumPy-named decompositions.

matmul computes op(a) @ op(b) in float32 under the profile mojolearn.identical.gemm.fp32.v1: the same bits on NVIDIA and AMD for the same inputs. It supports the NN, NT and TN products; asking for both transposes raises ValueError. By default it requires the identical build and raises when another build is loaded; identical=False uses whichever build is loaded, with no identity guarantee. numeric_mode, profile and require_identical report and check what this process loaded.

matmul_bf16, matmul_int8 and matmul_int15 are the low-bit profiles, with to_bf16, from_bf16, quantize_int8, dequantize_int8, quantize_int15 and dequantize_int15 to prepare and recover their operands.

qr, eigh, svdvals and svd carry NumPy’s names and meanings for the modes they support (note that qr defaults to mode='r'); unsupported modes are refused by name rather than approximated. They run on the device on a GPU install and on the host binding on a CPU-only install, with the same bits. Cholesky factors a symmetric positive-definite matrix and solves with the factor.

Importing this module does not load the extension or run any computation.

Examples

>>> import numpy as np
>>> from mojolearn import linalg
>>> a = np.random.default_rng(0).random((256, 64), dtype=np.float32)
>>> b = np.random.default_rng(1).random((64, 32), dtype=np.float32)
>>> c = np.asarray(linalg.matmul(a, b))

Functions

matmul(a, b, *[, transpose_a, transpose_b, ...])	Float32 matrix product op(a) @ op(b) with the same bits on NVIDIA and AMD.
numeric_mode()	The numeric mode of the linalg extension this process loaded.
profile()	Describe the matrix-product profile this process loaded, as a dict.
require_identical()	Raise RuntimeError unless this process loaded the identical build.
qr(a[, mode])	QR decomposition, as numpy.linalg.qr.
eigh(a[, UPLO])	Eigenvalues and eigenvectors of a symmetric matrix, as numpy.linalg.eigh.
svdvals(a)	Singular values in descending order, as numpy.linalg.svdvals.
svd(a[, full_matrices, compute_uv, hermitian])	Singular value decomposition, as numpy.linalg.svd.
matmul_bf16(a, b, *[, transpose_a, ...])	Matrix product with bfloat16 operands, under the profile mojolearn.identical.gemm.bf16f32.v1.
matmul_int8(a, b, *[, out])	Matrix product a @ b.T with int8 operands, under the profile mojolearn.identical.gemm.int8i32.v1.
to_bf16(x)	Convert float32 to bfloat16 bits.
from_bf16(bits)	Convert bfloat16 bits (a uint16 buffer) to float32, exactly.
quantize_int8(x)	Quantize a 2-D float32 matrix to int8 codes with one power-of-two exponent per row.
dequantize_int8(codes, exponents)	Recover float32 values codes * 2**exponents[row] exactly from quantize_int8 output.
matmul_int15(a, b, *[, out])	Matrix product a @ b.T with fifteen-bit operands, under the profile mojolearn.identical.gemm.int15i64.v1.
quantize_int15(x)	Quantize a 2-D float32 matrix to fifteen-bit codes stored as two int8 planes.
dequantize_int15(hi, lo, exponents)	Recover float32 values (hi * 128 + lo) * 2**exponents[row] from quantize_int15 output.

Classes

Cholesky([jitter])

Cholesky factorization A + jitter * I = L @ L.T of a symmetric positive-definite matrix, then solves.

mojolearn.linalg.matmul(a, b, *, transpose_a=False, transpose_b=False, out=None, identical=True)

Float32 matrix product op(a) @ op(b) with the same bits on NVIDIA and AMD.

Computes under the profile mojolearn.identical.gemm.fp32.v1. Each output cell accumulates its k products in a fixed order that depends only on the input bits, k and the profile: not on m, n, the launch geometry, the batch or the vendor. The order is chosen for reproducibility, not accuracy, so the result does not match NumPy’s a @ b bit for bit.

Parameters:

• a, b (float32 buffer, 2-D) – Any object exporting a 2-D float32 buffer: a NumPy array, an array.array('f'), a mojolearn.Array, and so on. NumPy is not required. Other dtypes raise TypeError rather than being cast. Non-contiguous input is copied; C-contiguous input is used without a copy.

• transpose_a, transpose_b (bool, default False) – Whether the stored array is the transpose of the operand. Three combinations are supported:

  transpose_a	transpose_b	C	a.shape	b.shape
  False	False	a @ b	(m, k)	(k, n)
  False	True	a @ b.T	(m, k)	(n, k)
  True	False	a.T @ b	(k, m)	(k, n)

  Both True raises ValueError; compute (b @ a).T instead, for example np.ascontiguousarray(np.asarray(matmul(b, a)).T).

• out (writable float32 buffer, optional) – Destination, C-contiguous with shape (m, n). A new mojolearn.Array is allocated when None.

• identical (bool, default True) – Require the identical build: raises RuntimeError when another build is loaded. False uses whichever build is loaded, with no identity guarantee.

Returns:

mojolearn.Array or out – The (m, n) float32 result; numpy.asarray on it is zero-copy. When out is given, that same object is returned.

Raises:

• ValueError – Both transpose flags set, mismatched contracted extents, or a wrong out shape or layout.

• TypeError – A non-float32 operand or out.

Notes

Compare NaN cells with “is NaN”, not by bits: NaN payloads are not guaranteed. A min, max or argmin you take over the output can depend on order, because -0.0 == +0.0; make such reductions order-independent yourself.

Examples

>>> import numpy as np
>>> from mojolearn import linalg
>>> a = np.random.default_rng(0).random((128, 64), dtype=np.float32)
>>> b = np.random.default_rng(1).random((32, 64), dtype=np.float32)
>>> c = np.asarray(linalg.matmul(a, b, transpose_b=True))

mojolearn.linalg.numeric_mode()

The numeric mode of the linalg extension this process loaded.

Returns ‘identical’, ‘deterministic’ or ‘fast’, read back from the loaded binary itself rather than from the environment, and checked against the directory it was loaded from and the binary’s profile version.

Returns:

str

Raises:

RuntimeError – When the binary is in the wrong directory, its profile version is unknown, or the extension is stale.

mojolearn.linalg.profile()

Describe the matrix-product profile this process loaded, as a dict.

Keys include profile (the profile name, for example 'mojolearn.identical.gemm.fp32.v1'), profile_version, numeric_mode, identity_claimed, dtype, ops and binary (the loaded file). Quote the profile name with any result you publish.

identity_claimed is True only when the identical build is loaded; in any other mode the remaining fields describe the computation but guarantee nothing about the bits.

Returns:

dict

mojolearn.linalg.require_identical()

Raise RuntimeError unless this process loaded the identical build.

Call it once at start-up when your program depends on bit-identical results, so a wrong build fails immediately rather than at the first product. matmul calls it for you unless you pass identical=False.

mojolearn.linalg.qr(a, mode='r')

QR decomposition, as numpy.linalg.qr.

Parameters:

• a (float32 buffer, shape (M, N)) – Non-float32 input raises TypeError rather than being cast.

• mode ({‘r’, ‘reduced’, ‘complete’, ‘raw’}, default ‘r’) – Unlike NumPy, the default is ‘r’ (R only); pass mode='reduced' for NumPy’s default (Q, R). ‘reduced’ returns Q (M, K) and R (K, N), ‘complete’ Q (M, M) and R (M, N), ‘raw’ h (N, M) and tau (K,), with K = min(M, N), following LAPACK’s conventions. Any other mode raises ValueError.

Returns:

• R (mojolearn.Array, shape (N, N), for mode=’r’) – Upper triangular, with R.T @ R equal to a.T @ a to float32 rounding. For M >= N its rows can differ in sign from qr(a, mode='reduced')[1].

• (Q, R) or (h, tau) – For the other modes.

Notes

Runs on the device on a GPU install and on the host binding on a CPU-only install, with the same bits.

mojolearn.linalg.eigh(a, UPLO='L')

Eigenvalues and eigenvectors of a symmetric matrix, as numpy.linalg.eigh.

Parameters:

• a (float32 buffer, shape (N, N)) – Read as symmetric from one triangle; the other is never read. A non-square matrix raises ValueError.

• UPLO ({‘L’, ‘U’}, default ‘L’) – Which triangle to read.

Returns:

• w (mojolearn.Array, shape (N,)) – Eigenvalues in ascending order.

• v (mojolearn.Array, shape (N, N)) – v[:, i] is the unit eigenvector of w[i]. Each vector’s sign is fixed by a deterministic rule, so the result depends only on the matrix.

Raises:

RuntimeError – When the Jacobi iteration does not converge, or the device kernel did not run.

Notes

Uses a round-robin Jacobi method, on the device on a GPU install and on the host binding on a CPU-only install, with the same bits.

mojolearn.linalg.svdvals(a)

Singular values in descending order, as numpy.linalg.svdvals.

Computed as the singular values of R from a Householder QR of a, which equal those of a; a wide matrix (M < N) is read through its transpose.

Parameters:

a (float32 buffer, shape (M, N))

Returns:

mojolearn.Array, shape (min(M, N),)

mojolearn.linalg.svd(a, full_matrices=True, compute_uv=True, hermitian=False)

Singular value decomposition, as numpy.linalg.svd.

Returns SVDResult(U, S, Vh) with S descending. The signs of the columns of U and V are deterministic but need not match LAPACK’s (NumPy promises no sign either). compute_uv=False returns svdvals. hermitian=True goes through eigh (lower triangle) as NumPy does and needs a square matrix, otherwise ValueError.

Parameters:

• a (float32 buffer, shape (M, N))

• full_matrices (bool, default True) – Return U (M, M) and Vh (N, N) instead of the reduced shapes.

• compute_uv (bool, default True) – Return U and Vh as well as S.

• hermitian (bool, default False) – Treat a as symmetric.

Returns:

SVDResult or mojolearn.Array

class mojolearn.linalg.Cholesky(jitter=None)

Bases: NumericModeMixin

Cholesky factorization A + jitter * I = L @ L.T of a symmetric positive-definite matrix, then solves.

fit factors a given float32 matrix once; solve answers A X = B from the factor and logdet_ gives the log-determinant. A failed factorization is not an exception: check info_ (0 means success); a failed factor’s solve and logdet_ refuse. The factor can be saved with save and loaded with load, including on a CPU-only install, where fit and solve also run.

Parameters:

jitter (float or None, default None) – Ridge added to the diagonal before factoring. None uses the profile’s fixed ridge (see profile_jitter); 0.0 adds none. Any other value is refused, because the identical mode fixes the ridge to those two values. A non-number raises TypeError.

Variables:

• L_ (mojolearn.Array, shape (n, n), float32) – Lower-triangular factor, strict upper triangle zero. Meaningful only when info_ == 0.

• info_ (int) – LAPACK-style status: 0 for a complete factor; k > 0 means the leading minor of order k is not positive definite and L_ is partial.

• nb_ (int) – Panel width used by the factorization.

• logdet_ (float) – log |A + jitter * I|. Raises ValueError after a failed fit.

• jitter_ (float) – The ridge that was added.

• n_ (int) – Matrix size.

Notes

A non-square A or a right-hand side whose row count is not n raises ValueError; non-finite or non-symmetric matrices are refused before any work.

Examples

>>> import numpy as np, mojolearn
>>> rng = np.random.default_rng(0)
>>> M = rng.random((64, 64), dtype=np.float32)
>>> A = (M @ M.T + 64 * np.eye(64, dtype=np.float32)).astype(np.float32)
>>> chol = mojolearn.Cholesky(jitter=0.0).fit(A)
>>> x = np.asarray(chol.solve(np.ones(64, dtype=np.float32)))

vendor_used()

Accelerator API of the loaded binary, read back from it; “cpu” on a CPU-only install.

profile_jitter()

The profile’s fixed ridge, used when jitter=None.

fit(A)

Factor A + jitter * I; returns self.

A is a square 2-D float32 matrix, otherwise ValueError. A failed factorization sets info_ != 0 instead of raising.

property logdet_

log |A + jitter * I|; raises ValueError after a failed fit.

solve(B)

Solve A X = B from the factor.

Parameters:

B (float32 array, shape (n,) or (n, nrhs))

Returns:

X (float32 array, the shape of B)

Raises:

ValueError – Before fit, or when B has the wrong row count or no columns. Solving against a failed factor is also refused.

save(path)

Write the factor to path as an .npz file.

Stores the factor, info_, nb_, logdet_, the jitter and the numeric mode. A failed factor saves too; its solve still refuses. Raises RuntimeError before fit.

classmethod load(path)

Load a factor written by save.

The result can solve and report logdet_; it holds no matrix and does not refactor. A file saved by another estimator, or with mismatched shapes or dtypes, raises ValueError.

mojolearn.linalg.matmul_bf16(a, b, *, transpose_a=False, transpose_b=False, out=None)

Matrix product with bfloat16 operands, under the profile mojolearn.identical.gemm.bf16f32.v1.

b is bfloat16 bits (a uint16 buffer from to_bf16); a is float32 or bfloat16 bits too. The output is float32: the operands are widened exactly and multiplied with matmul’s arithmetic, so matmul(from_bf16(a_bits), from_bf16(b_bits)) gives the same bits. transpose_a, transpose_b and out behave as in matmul. Requires the identical build.

mojolearn.linalg.matmul_int8(a, b, *, out=None)

Matrix product a @ b.T with int8 operands, under the profile mojolearn.identical.gemm.int8i32.v1.

Each of a (m, k) and b (n, k) is either a float32 matrix, quantized by the profile’s rule, or a (codes, exponents) pair from quantize_int8. The output is float32: the exact int32 sum of the code products, scaled by a power of two. It is the product of the rounded operands, not the float32 product. out behaves as in matmul. Requires the identical build.

mojolearn.linalg.to_bf16(x)

Convert float32 to bfloat16 bits.

Subnormals are flushed, then each value is rounded to nearest even. Returns a uint16 mojolearn.Array of the same shape, the input format of matmul_bf16.

mojolearn.linalg.from_bf16(bits)

Convert bfloat16 bits (a uint16 buffer) to float32, exactly.

Returns a float32 mojolearn.Array of the same shape.

mojolearn.linalg.quantize_int8(x)

Quantize a 2-D float32 matrix to int8 codes with one power-of-two exponent per row.

Returns:

• codes (mojolearn.Array of int8) – The shape of x.

• exponents (mojolearn.Array of int32) – One per row; row = codes * 2**exponent.

mojolearn.linalg.dequantize_int8(codes, exponents)

Recover float32 values codes * 2**exponents[row] exactly from quantize_int8 output.

mojolearn.linalg.matmul_int15(a, b, *, out=None)

Matrix product a @ b.T with fifteen-bit operands, under the profile mojolearn.identical.gemm.int15i64.v1.

Each of a (m, k) and b (n, k) is either a float32 matrix, quantized by the profile’s rule, or a (hi, lo, exponents) triple from quantize_int15. The output is float32: the exact integer sum of the code products, rounded to nearest even and scaled by a power of two. It is the product of the rounded operands, not the float32 product. k may be at most 65536 (INT15_MAX_K); larger values raise ValueError. out behaves as in matmul. Requires the identical build.

mojolearn.linalg.quantize_int15(x)

Quantize a 2-D float32 matrix to fifteen-bit codes stored as two int8 planes.

Per row, e = floor(log2(absmax)) - 13 so the largest magnitude lands in [8192, 16384); each code is clamp(round_half_even(x * 2**-e), -16383, 16383), split as hi = code >> 7 and lo = code & 127 so code = hi * 128 + lo. The same planes and exponents on every vendor.

Returns:

• hi, lo (mojolearn.Array of int8) – The shape of x.

• exponents (mojolearn.Array of int32) – One per row.

mojolearn.linalg.dequantize_int15(hi, lo, exponents)

Recover float32 values (hi * 128 + lo) * 2**exponents[row] from quantize_int15 output.