Repository navigation
Einsum selects unsupported matmul for integer contractions #4463
Description
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.
Activity
einsum: Fallback to multiply-and-reduce for integer contractions
zcbenz commentedon Sep 5, 2026
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 commentedon Sep 5, 2026
@zcbenz We can add a note here something like this:
mlx/python/src/ops.cpp
Lines 5782 to 5787 in 2d27ab0
Also I saw we can add similar note here as well.
mlx/python/src/ops.cpp
Lines 732 to 748 in 2d27ab0
Let me know if we need to document this somewhere else as well.
simeetnayan81 commentedon Sep 5, 2026
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 commentedon Sep 5, 2026
@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 commentedon Sep 6, 2026
Thanks, will raise a PR. Or if someone else is interested please go ahead.
zcbenz commentedon Sep 30, 2026
Closed by #4550