Einsum selects unsupported matmul for integer contractions
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 28.5k
- Forks
- 2.3k
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 62
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the implementation behind mx.einsum and trace how the "ij,jk->ik" contraction selects the matmul path. Run the provided int32 reproduction and inspect existing multiply-and-reduce Einsum behavior for a compatible fallback. Done means integer contractions produce the expected result without matmul's floating-point dtype error, with relevant tests added or updated.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100