diff --git a/mlx/backend/common/common.cpp b/mlx/backend/common/common.cpp index cbc90ed27e..5426301a35 100644 --- a/mlx/backend/common/common.cpp +++ b/mlx/backend/common/common.cpp @@ -3,6 +3,7 @@ #include "mlx/backend/common/broadcasting.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -99,50 +100,10 @@ void NumberOfElements::eval(const std::vector& inputs, array& out) { numel = 1.0 / numel; } - switch (out.dtype()) { - case bool_: - *out.data() = static_cast(numel); - break; - case uint8: - *out.data() = static_cast(numel); - break; - case uint16: - *out.data() = static_cast(numel); - break; - case uint32: - *out.data() = static_cast(numel); - break; - case uint64: - *out.data() = static_cast(numel); - break; - case int8: - *out.data() = static_cast(numel); - break; - case int16: - *out.data() = static_cast(numel); - break; - case int32: - *out.data() = static_cast(numel); - break; - case int64: - *out.data() = static_cast(numel); - break; - case float16: - *out.data() = static_cast(numel); - break; - case float32: - *out.data() = static_cast(numel); - break; - case bfloat16: - *out.data() = static_cast(numel); - break; - case float64: - *out.data() = static_cast(numel); - break; - case complex64: - *out.data() = static_cast(numel); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + *out.data() = static_cast(numel); + }); } std::pair prepare_reshape(const array& in, const array& out) {