ml-explore / ml-explore/mlx

Einsum selects unsupported matmul for integer contractions

Open
#4,463 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

documentation low priority
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.