| 48 | } |
| 49 | |
| 50 | std::pair<mx::array, mx::array> to_arrays( |
| 51 | const ScalarOrArray& a, |
| 52 | const ScalarOrArray& b) { |
| 53 | // Four cases: |
| 54 | // - If both a and b are arrays leave their types alone |
| 55 | // - If a is an array but b is not, treat b as a weak python type |
| 56 | // - If b is an array but a is not, treat a as a weak python type |
| 57 | // - If neither is an array convert to arrays but leave their types alone |
| 58 | auto is_mlx_array = [](const ScalarOrArray& x) { |
| 59 | return std::holds_alternative<mx::array>(x) || |
| 60 | std::holds_alternative<ArrayLike>(x) && |
| 61 | nb::hasattr(std::get<ArrayLike>(x).obj, "__mlx_array__"); |
| 62 | }; |
| 63 | auto get_mlx_array = [](const ScalarOrArray& x) { |
| 64 | if (auto px = std::get_if<mx::array>(&x); px) { |
| 65 | return *px; |
| 66 | } else { |
| 67 | return nb::cast<mx::array>( |
| 68 | std::get<ArrayLike>(x).obj.attr("__mlx_array__")); |
| 69 | } |
| 70 | }; |
| 71 | |
| 72 | if (is_mlx_array(a)) { |
| 73 | auto arr_a = get_mlx_array(a); |
| 74 | if (is_mlx_array(b)) { |
| 75 | auto arr_b = get_mlx_array(b); |
| 76 | return {arr_a, arr_b}; |
| 77 | } |
| 78 | return {arr_a, to_array(b, arr_a.dtype())}; |
| 79 | } else if (is_mlx_array(b)) { |
| 80 | auto arr_b = get_mlx_array(b); |
| 81 | return {to_array(a, arr_b.dtype()), arr_b}; |
| 82 | } else { |
| 83 | return {to_array(a), to_array(b)}; |
| 84 | } |
| 85 | } |
| 86 | |
| 87 | mx::array to_array_with_accessor(nb::object obj) { |
| 88 | if (nb::isinstance<mx::array>(obj)) { |