mojolearn.matmul#

mojolearn.matmul(a, b, *, transpose_a=False, transpose_b=False, out=None, identical=True)[source]#

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))
On this page