| 447 | }; |
| 448 | |
| 449 | nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) { |
| 450 | // Flat array inputs |
| 451 | std::vector<mx::array> inputs; |
| 452 | |
| 453 | // Compilation constants which includes the tree structure of the arguments |
| 454 | std::vector<uint64_t> constants; |
| 455 | |
| 456 | // Reserve some large primes to signify the presence of an array, a list or |
| 457 | // a dict in order to encode the structure of the pytree. We choose primes |
| 458 | // to reduce slightly the chances of these numbers occurring by a |
| 459 | // multiplication as values in the constants list. |
| 460 | constexpr uint64_t array_identifier = 18446744073709551557UL; |
| 461 | constexpr uint64_t list_identifier = 18446744073709551533UL; |
| 462 | constexpr uint64_t dict_identifier = 18446744073709551521UL; |
| 463 | constexpr uint64_t none_identifier = 10239356951478402889UL; |
| 464 | |
| 465 | // Flatten the tree with hashed constants and structure |
| 466 | std::function<void(nb::handle)> recurse; |
| 467 | recurse = [&](nb::handle obj) { |
| 468 | if (nb::isinstance<nb::list>(obj)) { |
| 469 | auto l = nb::cast<nb::list>(obj); |
| 470 | constants.push_back(list_identifier); |
| 471 | for (int i = 0; i < l.size(); ++i) { |
| 472 | recurse(l[i]); |
| 473 | } |
| 474 | } else if (nb::isinstance<nb::tuple>(obj)) { |
| 475 | auto l = nb::cast<nb::tuple>(obj); |
| 476 | constants.push_back(list_identifier); |
| 477 | for (auto item : obj) { |
| 478 | recurse(item); |
| 479 | } |
| 480 | } else if (nb::isinstance<nb::dict>(obj)) { |
| 481 | auto d = nb::cast<nb::dict>(obj); |
| 482 | constants.push_back(dict_identifier); |
| 483 | for (auto item : d) { |
| 484 | auto r = item.first.attr("__hash__")(); |
| 485 | constants.push_back(nb::cast<int64_t>(r)); |
| 486 | recurse(item.second); |
| 487 | } |
| 488 | } else if (nb::isinstance<mx::array>(obj)) { |
| 489 | inputs.push_back(nb::cast<mx::array>(obj)); |
| 490 | constants.push_back(array_identifier); |
| 491 | } else if (nb::isinstance<nb::str>(obj)) { |
| 492 | auto r = obj.attr("__hash__")(); |
| 493 | constants.push_back(nb::cast<int64_t>(r)); |
| 494 | } else if (nb::isinstance<nb::int_>(obj)) { |
| 495 | constants.push_back(nb::cast<int64_t>(obj)); |
| 496 | } else if (nb::isinstance<nb::float_>(obj)) { |
| 497 | auto r = nb::cast<double>(obj); |
| 498 | constants.push_back(*reinterpret_cast<uint64_t*>(&r)); |
| 499 | } else if (obj.is_none()) { |
| 500 | constants.push_back(none_identifier); |
| 501 | } else { |
| 502 | std::ostringstream msg; |
| 503 | msg << "[compile] Function arguments must be trees of arrays " |
| 504 | << "or constants (floats, ints, strings, or None), but received " |
| 505 | << "type " << type_name_str(obj) << "."; |
| 506 | throw std::invalid_argument(msg.str()); |
no test coverage detected