Contract integer inputs in float for dot, tensordot and inner - #24
Open
MarcosAsh wants to merge 1 commit into
Open
Contract integer inputs in float for dot, tensordot and inner#24MarcosAsh wants to merge 1 commit into
MarcosAsh wants to merge 1 commit into
Conversation
SamanehSaadat
approved these changes
Aug 21, 2026
MarcosAsh
force-pushed
the
split-int-contractions
branch
from
August 30, 2026 11:45
4b602ae to
937cdc3
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
matmulalready computed integer contractions in float32 and cast back, since mlx has no integer contraction kernel.dot,tensordotandinnerdid not, they went straight to@,mx.tensordotandmx.inner, so every integer combination raised[matmul] Only inexact types are supported.The dtype choice is now a
_contraction_dtypeshelper that all five contraction ops share,einsumincluded, which had its own copy of the same lines.matmulkeeps the int8 to int32 accumulation rule, which is specific to it and comes from keras passingpreferred_element_typein 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
matmulhas been making, and mlx has no integer kernel to fall back on.