Einsum selects unsupported matmul for integer contractions

[justinchuby]

Problem

mx.einsum supports integer inputs for non-matmul forms, but a valid integer contraction such as "ij,jk->ik" is lowered to matmul, which rejects non-floating types.

This is related to #516, which declined dedicated integer GEMM kernels. einsum may still be able to support these equations without an integer GEMM by using its existing multiply-and-reduce lowering as a fallback when the optimized matmul path does not support the dtype.

This matters for ONNX interoperability because ONNX Einsum-28 permits all numeric tensor types. In mlx-c applications, reaching this error with the default error handler terminates the host process (exit(-1)), so downstream runtimes must currently reject integer contractions before calling MLX.

Reproduction

import mlx.core as mx

x = mx.array([[1, 2], [3, 4]], dtype=mx.int32)

# Integer Einsum forms without a matmul contraction work.
print(mx.einsum("ij->ji", x))
print(mx.einsum("ij->", x))
print(mx.einsum("i,j->ij", x[0], x[1]))

# A contraction selected for the matmul path fails.
y = mx.einsum("ij,jk->ik", x, x)
mx.eval(y)

Observed with MLX 0.32.2:

ValueError: [matmul] Only inexact types are supported but int32 and int32 were provided which results in int32, which is not a floating point type.

The same behavior occurs for int8/int16/int64 and uint8/uint16/uint32/uint64.

Expected behavior

Preferably, integer contractions should fall back to an integer-compatible multiply-and-reduce Einsum path when the optimized matmul path is unavailable. If integer contractions are intentionally unsupported, einsum should document that only non-contraction integer equations are supported.

[zcbenz]

While we can redirect to a native path I think it would be unexpected when someone actually wants to do integer contraction. Integer matmul was not supported as I remember was because it was decided that it was not useful in practice, so unless there is a real use case, I prefer to just document the behavior.

[simeetnayan81]

@zcbenz We can add a note here something like this:

Note: Contractions require inexact types (floating or complex),
the same as :func:`matmul`. Integer and boolean arrays are
supported only for equations that do not contract an axis.

mlx/python/src/ops.cpp

Lines 5782 to 5787 in 2d27ab0

	nb::sig("def einsum_path(subscripts: str, *operands)"),
	R"pbdoc(
	Compute the contraction order for the given Einstein summation.
	Args:

Also I saw we can add similar note here as well.

mlx/python/src/ops.cpp

Lines 732 to 748 in 2d27ab0

	"def matmul(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
	R"pbdoc(
	Matrix multiplication.
	Perform the (possibly batched) matrix multiplication of two arrays. This function supports
	broadcasting for arrays with more than two dimensions.
	- If the first array is 1-D then a 1 is prepended to its shape to make it
	a matrix. Similarly if the second array is 1-D then a 1 is appended to its
	shape to make it a matrix. In either case the singleton dimension is removed
	from the result.
	- A batched matrix multiplication is performed if the arrays have more than
	2 dimensions. The matrix dimensions for the matrix product are the last
	two dimensions of each input.
	- All but the last two dimensions of each input are broadcast with one another using
	standard numpy-style broadcasting semantics.

Let me know if we need to document this somewhere else as well.

[simeetnayan81]

Preferably, integer contractions should fall back to an integer-compatible multiply-and-reduce Einsum path when the optimized matmul path is unavailable.

I feel if we are doing this then we should do this consistently for other operations as well like matmul, tensordot, inner, @, gather_mm etc. Until we have plans for consistently supporting these.

[zcbenz]

@simeetnayan81 I think we can say something simpler like "Contractions are implemented with matmul and have same requirements", and then add documentations to matmul.

[simeetnayan81]

Thanks, will raise a PR. Or if someone else is interested please go ahead.