[BUG] mlx crashes with msg - uncaught exception of type std::invalid_argument: [Scatter::eval_gpu] Does not support int64

[lkarthee]

[BUG] mlx crashes with msg - uncaught exception of type std::invalid_argument: [Scatter::eval_gpu] Does not support int64

[awni]

If you are just asking for a catchable exception then #1077 should close this. We would like to eventually allow int64 and other 8 byte types to work with scatter, but that is more involved.

[lkarthee]

Thank you Awni. some observations:

• crash message is so confusing - does not say where the problem is with the array or indices or values. Can we improve it by mentioning workaround in error message ?
• scatter ops can use cpu device for int64 and uint64 ?
• adding a note about supported dtypes and devices cpu, gpu in mlx.core.array.at is helpful .
• are there any other ops which are not supported on gpu and run on cpu ?

zeros = mx.zeros(shape, dtype=values.dtype) 
zeros = zeros.at[indices].add(values) 

i tried this and it does not work as add does not take device kw_arg:

if zeros.dtype in [mx.int64, mx.uint64] and mx.get_default_device == mx.DeviceType.gpu :
  device = mx.Device(type=mx.DeviceType.cpu)
  zeros = zeros.at[indices].add(values, device=device) 
else:
  zeros = zeros.at[indices].add(values)

It would be helpful if mlx can fallback to cpu for scatter ops which are not supported on gpu or allow device kw_arg for all scatter ops.

Additional ops which are impacted by this bug:

• mx.cumsum
• mx.cumprod
• mx.diag

[awni]

crash message is so confusing - does not say where the problem is with the array or indices or values. Can we improve it by mentioning workaround in error message ?

I improved the message in #1077. The problem is with the values.

scatter ops can use cpu device for int64 and uint64 ?

We prefer not to silently route to the CPU for ops without a GPU back-end. You can do this in the API by changing the default stream to the CPU before calling the scatter when the dytpe is int64/uint64.

are there any other ops which are not supported on gpu and run on cpu ?

Just a few. FFT and some of the lapack ops (QR / Inverse). Metal support for FFT is coming soon in #981 .

i tried this and it does not work as add does not take device kw_arg:

You can use a context manager. For most free ops stream kwarg also works. E.g.

v = mx.array([1, 2, 3])
u = mx.array([1, 2])
idx = mx.array([0, 1])

with mx.stream(mx.cpu):
    out = v.at[idx].add(u)

[lkarthee]

Thank you @awni for the fix.