Skip to content

Contract integer inputs in float for dot, tensordot and inner - #24

Open
MarcosAsh wants to merge 1 commit into
mainfrom
split-int-contractions
Open

Contract integer inputs in float for dot, tensordot and inner#24
MarcosAsh wants to merge 1 commit into
mainfrom
split-int-contractions

Conversation

@MarcosAsh

Copy link
Copy Markdown
Collaborator

Description

matmul already computed integer contractions in float32 and cast back, since mlx has no integer contraction kernel. dot, tensordot and inner did not, they went straight to @, mx.tensordot and mx.inner, so every integer combination raised [matmul] Only inexact types are supported.

The dtype choice is now a _contraction_dtypes helper that all five contraction ops share, einsum included, which had its own copy of the same lines. matmul keeps the int8 to int32 accumulation rule, which is specific to it and comes from keras passing preferred_element_type in its own test. The other three take plain promotion, which is what jax does for all three.

All 36 dot dtype pairs now match jax, and values match numpy exactly for int8, int16, int32, uint8 and uint16 across the three ops. This retires three exclusion entries covering 64 tests.

Worth knowing the limit, and the docstring says it: float32 is exact only to 2**24, and it is the products that have to fit rather than the inputs, so a contraction whose result passes that loses its low bits. That is the same trade matmul has been making, and mlx has no integer kernel to fall back on.

@MarcosAsh
MarcosAsh force-pushed the split-int-contractions branch from 4b602ae to 937cdc3 Compare August 30, 2026 11:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants