| 34 | } |
| 35 | |
| 36 | mx::array nd_array_to_mlx( |
| 37 | nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array, |
| 38 | std::optional<mx::Dtype> dtype, |
| 39 | std::optional<nb::dlpack::dtype> nb_dtype) { |
| 40 | // Compute the shape and size |
| 41 | mx::Shape shape; |
| 42 | shape.reserve(nd_array.ndim()); |
| 43 | for (int i = 0; i < nd_array.ndim(); i++) { |
| 44 | shape.push_back(check_shape_dim(nd_array.shape(i))); |
| 45 | } |
| 46 | auto type = nb_dtype.value_or(nd_array.dtype()); |
| 47 | |
| 48 | // Copy data and make array |
| 49 | if (type == nb::dtype<bool>()) { |
| 50 | return nd_array_to_mlx_contiguous<bool>( |
| 51 | nd_array, shape, dtype.value_or(mx::bool_)); |
| 52 | } else if (type == nb::dtype<uint8_t>()) { |
| 53 | return nd_array_to_mlx_contiguous<uint8_t>( |
| 54 | nd_array, shape, dtype.value_or(mx::uint8)); |
| 55 | } else if (type == nb::dtype<uint16_t>()) { |
| 56 | return nd_array_to_mlx_contiguous<uint16_t>( |
| 57 | nd_array, shape, dtype.value_or(mx::uint16)); |
| 58 | } else if (type == nb::dtype<uint32_t>()) { |
| 59 | return nd_array_to_mlx_contiguous<uint32_t>( |
| 60 | nd_array, shape, dtype.value_or(mx::uint32)); |
| 61 | } else if (type == nb::dtype<uint64_t>()) { |
| 62 | return nd_array_to_mlx_contiguous<uint64_t>( |
| 63 | nd_array, shape, dtype.value_or(mx::uint64)); |
| 64 | } else if (type == nb::dtype<int8_t>()) { |
| 65 | return nd_array_to_mlx_contiguous<int8_t>( |
| 66 | nd_array, shape, dtype.value_or(mx::int8)); |
| 67 | } else if (type == nb::dtype<int16_t>()) { |
| 68 | return nd_array_to_mlx_contiguous<int16_t>( |
| 69 | nd_array, shape, dtype.value_or(mx::int16)); |
| 70 | } else if (type == nb::dtype<int32_t>()) { |
| 71 | return nd_array_to_mlx_contiguous<int32_t>( |
| 72 | nd_array, shape, dtype.value_or(mx::int32)); |
| 73 | } else if (type == nb::dtype<int64_t>()) { |
| 74 | return nd_array_to_mlx_contiguous<int64_t>( |
| 75 | nd_array, shape, dtype.value_or(mx::int64)); |
| 76 | } else if (type == nb::dtype<mx::float16_t>()) { |
| 77 | return nd_array_to_mlx_contiguous<mx::float16_t>( |
| 78 | nd_array, shape, dtype.value_or(mx::float16)); |
| 79 | } else if (type == nb::dtype<mx::bfloat16_t>()) { |
| 80 | return nd_array_to_mlx_contiguous<mx::bfloat16_t>( |
| 81 | nd_array, shape, dtype.value_or(mx::bfloat16)); |
| 82 | } else if (type == nb::dtype<float>()) { |
| 83 | return nd_array_to_mlx_contiguous<float>( |
| 84 | nd_array, shape, dtype.value_or(mx::float32)); |
| 85 | } else if (type == nb::dtype<double>()) { |
| 86 | return nd_array_to_mlx_contiguous<double>( |
| 87 | nd_array, shape, dtype.value_or(mx::float32)); |
| 88 | } else if (type == nb::dtype<std::complex<float>>()) { |
| 89 | return nd_array_to_mlx_contiguous<mx::complex64_t>( |
| 90 | nd_array, shape, dtype.value_or(mx::complex64)); |
| 91 | } else if (type == nb::dtype<std::complex<double>>()) { |
| 92 | return nd_array_to_mlx_contiguous<mx::complex128_t>( |
| 93 | nd_array, shape, dtype.value_or(mx::complex64)); |
no test coverage detected