zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-04 16:21:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Allow scatter type exception to be caught by checking in op (#1077)

Browse Source 
* allow exception to be caught in main thread

* only for gpu

* more detailed scatter error
This commit is contained in:
Awni Hannun 2024-05-13 17:43:53 -07:00 committed by GitHub
parent 7178ac0111
commit 863039da4c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 54 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/backend/metal/kernels/reduction/reduce_inst.h
View File

@ -8,6 +8,7 @@
#include "mlx/backend/metal/kernels/defines.h" #include "mlx/backend/metal/kernels/defines.h"
#include "mlx/backend/metal/kernels/reduction/ops.h" #include "mlx/backend/metal/kernels/reduction/ops.h"
// clang-format off
#define instantiate_reduce_helper_floats(inst_f, name, op) \ #define instantiate_reduce_helper_floats(inst_f, name, op) \
inst_f(name, float16, half, op) inst_f(name, float32, float, op) \ inst_f(name, float16, half, op) inst_f(name, float32, float, op) \
inst_f(name, bfloat16, bfloat16_t, op) inst_f(name, bfloat16, bfloat16_t, op)
@ -38,7 +39,8 @@
inst_f(name##tname, itype, otype, op) inst_f(name##tname, itype, otype, op)
#define instantiate_reduce_from_types(inst_f, name, otype, op) \ #define instantiate_reduce_from_types(inst_f, name, otype, op) \
instantiate_reduce_from_types_helper(inst_f, name, bool_, bool, otype, op) \ instantiate_reduce_from_types_helper( \
inst_f, name, bool_, bool, otype, op) \
instantiate_reduce_from_types_helper( \ instantiate_reduce_from_types_helper( \
inst_f, name, uint8, uint8_t, otype, op) \ inst_f, name, uint8, uint8_t, otype, op) \
instantiate_reduce_from_types_helper( \ instantiate_reduce_from_types_helper( \
@ -69,3 +71,4 @@
bfloat16_t, \ bfloat16_t, \
otype, \ otype, \
op) op)
// clang-format on

8
mlx/ops.cpp
View File

@ -2641,6 +2641,14 @@ array scatter(
idx = astype(idx, dtype, s); idx = astype(idx, dtype, s);
} }
// TODO, remove when scatter supports 64-bit outputs
if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8) {
std::ostringstream msg;
msg << "[scatter] GPU scatter does not yet support " << a.dtype()
<< " for the input or updates.";
throw std::invalid_argument(msg.str());
}
inputs.insert(inputs.begin(), a); inputs.insert(inputs.begin(), a);
// TODO promote or cast? // TODO promote or cast?
inputs.push_back(astype(updates, a.dtype(), s)); inputs.push_back(astype(updates, a.dtype(), s));