What happens
diag on a 1-D input fails on the GPU for any 8-byte dtype, and works on the CPU:
import mlx.core as mx
mx.eval(mx.diag(mx.array([1, 2, 3], dtype=mx.int64)))
# [scatter] GPU scatter does not yet support int64 for the input or updates
uint64 and complex64 fail the same way. int32 and float32 are fine, and every
one of them works on the CPU stream.
Why
diag places the values with a scatter (mlx/ops.cpp:6371), and Scatter::eval_gpu
rejects any 8-byte output that is not a complex64 sum
(mlx/backend/metal/indexing.cpp:233). Nothing above that guard knows the restriction
exists, so an op that has nothing to do with scatter fails on a dtype it otherwise
supports.
eye had the same problem through the same path and was fixed in #4301 by building in a
narrower type and casting. That is not available here. eye only ever writes 0 and 1,
which are exact everywhere; these are the caller's values, and a cast through float32
changes them:
in 9007199254740993 4611686018427387903
via float32 9007199254740992 4611686018427387904
Environment
mlx built from c2bcf47, M5 Max and M4 Pro, macOS 26.5. Also reproduced on top of #4301,
so it is independent of that fix.
Fix
PR below. The 8-byte types select the diagonal out of a broadcast instead of scattering
into it, which never leaves the dtype. The existing scatter path is untouched for every
type that works today, since the mask costs an n x n allocation the scatter does not.
What happens
diagon a 1-D input fails on the GPU for any 8-byte dtype, and works on the CPU:uint64andcomplex64fail the same way.int32andfloat32are fine, and everyone of them works on the CPU stream.
Why
diagplaces the values with ascatter(mlx/ops.cpp:6371), andScatter::eval_gpurejects any 8-byte output that is not a
complex64sum(
mlx/backend/metal/indexing.cpp:233). Nothing above that guard knows the restrictionexists, so an op that has nothing to do with scatter fails on a dtype it otherwise
supports.
eyehad the same problem through the same path and was fixed in #4301 by building in anarrower type and casting. That is not available here.
eyeonly ever writes 0 and 1,which are exact everywhere; these are the caller's values, and a cast through float32
changes them:
Environment
mlx built from c2bcf47, M5 Max and M4 Pro, macOS 26.5. Also reproduced on top of #4301,
so it is independent of that fix.
Fix
PR below. The 8-byte types select the diagonal out of a broadcast instead of scattering
into it, which never leaves the dtype. The existing scatter path is untouched for every
type that works today, since the mask costs an n x n allocation the scatter does not.