| 458 | } |
| 459 | |
| 460 | mx::array create_array(nb::object v, std::optional<mx::Dtype> t) { |
| 461 | if (nb::isinstance<nb::bool_>(v)) { |
| 462 | return mx::array(nb::cast<bool>(v), t.value_or(mx::bool_)); |
| 463 | } else if (nb::isinstance<nb::int_>(v)) { |
| 464 | auto val = nb::cast<int64_t>(v); |
| 465 | auto default_type = (val > std::numeric_limits<int>::max() || |
| 466 | val < std::numeric_limits<int>::min()) |
| 467 | ? mx::int64 |
| 468 | : mx::int32; |
| 469 | return mx::array(val, t.value_or(default_type)); |
| 470 | } else if (nb::isinstance<nb::float_>(v)) { |
| 471 | auto out_type = t.value_or(mx::float32); |
| 472 | if (out_type == mx::float64) { |
| 473 | return mx::array(nb::cast<double>(v), out_type); |
| 474 | } else { |
| 475 | return mx::array(nb::cast<float>(v), out_type); |
| 476 | } |
| 477 | } else if (PyComplex_Check(v.ptr())) { |
| 478 | return mx::array( |
| 479 | static_cast<mx::complex64_t>(nb::cast<std::complex<float>>(v)), |
| 480 | t.value_or(mx::complex64)); |
| 481 | } else if (nb::isinstance<nb::list>(v)) { |
| 482 | return array_from_list(nb::cast<nb::list>(v), t); |
| 483 | } else if (nb::isinstance<nb::tuple>(v)) { |
| 484 | return array_from_list(nb::cast<nb::tuple>(v), t); |
| 485 | } else if (nb::isinstance<mx::array>(v)) { |
| 486 | auto arr = nb::cast<mx::array>(v); |
| 487 | return mx::astype(arr, t.value_or(arr.dtype())); |
| 488 | } else if (nb::ndarray_check(v)) { |
| 489 | using ContigArray = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>; |
| 490 | ContigArray nd; |
| 491 | std::optional<nb::dlpack::dtype> nb_dtype; |
| 492 | // Nanobind does not recognize bfloat16 numpy array: |
| 493 | // https://github.com/wjakob/nanobind/discussions/560 |
| 494 | if (nb::hasattr(v, "dtype") && v.attr("dtype").equal(nb::str("bfloat16"))) { |
| 495 | nd = nb::cast<ContigArray>(v.attr("view")("uint16")); |
| 496 | nb_dtype = nb::dtype<mx::bfloat16_t>(); |
| 497 | } else { |
| 498 | nd = nb::cast<ContigArray>(v); |
| 499 | } |
| 500 | return nd_array_to_mlx(nd, t, nb_dtype); |
| 501 | } else { |
| 502 | auto arr = to_array_with_accessor(v); |
| 503 | return mx::astype(arr, t.value_or(arr.dtype())); |
| 504 | } |
| 505 | } |
no test coverage detected