| 6 | #include "python/src/convert.h" |
| 7 | |
| 8 | mx::array to_array( |
| 9 | const ScalarOrArray& v, |
| 10 | std::optional<mx::Dtype> dtype /* = std::nullopt */) { |
| 11 | if (auto pv = std::get_if<nb::bool_>(&v); pv) { |
| 12 | return mx::array(nb::cast<bool>(*pv), dtype.value_or(mx::bool_)); |
| 13 | } else if (auto pv = std::get_if<nb::int_>(&v); pv) { |
| 14 | auto val = nb::cast<int64_t>(*pv); |
| 15 | auto default_type = (val > std::numeric_limits<int>::max() || |
| 16 | val < std::numeric_limits<int>::min()) |
| 17 | ? mx::int64 |
| 18 | : mx::int32; |
| 19 | auto out_t = dtype.value_or(default_type); |
| 20 | if (mx::issubdtype(out_t, mx::integer) && out_t.size() < 8) { |
| 21 | auto info = mx::iinfo(out_t); |
| 22 | if (val < info.min || val > static_cast<int64_t>(info.max)) { |
| 23 | std::ostringstream msg; |
| 24 | msg << "Converting " << val << " to " << out_t |
| 25 | << " would result in overflow."; |
| 26 | throw std::invalid_argument(msg.str()); |
| 27 | } |
| 28 | } |
| 29 | |
| 30 | // bool_ is an exception and is always promoted |
| 31 | return mx::array(val, (out_t == mx::bool_) ? mx::int32 : out_t); |
| 32 | } else if (auto pv = std::get_if<nb::float_>(&v); pv) { |
| 33 | auto out_t = dtype.value_or(mx::float32); |
| 34 | return mx::array( |
| 35 | nb::cast<float>(*pv), |
| 36 | mx::issubdtype(out_t, mx::floating) ? out_t : mx::float32); |
| 37 | } else if (auto pv = std::get_if<std::complex<float>>(&v); pv) { |
| 38 | return mx::array(static_cast<mx::complex64_t>(*pv), mx::complex64); |
| 39 | } else if (auto pv = std::get_if<mx::array>(&v); pv) { |
| 40 | return *pv; |
| 41 | } else if (auto pv = std::get_if< |
| 42 | nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>>(&v); |
| 43 | pv) { |
| 44 | return nd_array_to_mlx(*pv, dtype); |
| 45 | } else { |
| 46 | return to_array_with_accessor(std::get<ArrayLike>(v).obj); |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | std::pair<mx::array, mx::array> to_arrays( |
| 51 | const ScalarOrArray& a, |
no test coverage detected