| 838 | } |
| 839 | |
| 840 | std::pair<std::vector<mx::array>, std::vector<int>> operator()( |
| 841 | const std::vector<mx::array>& inputs, |
| 842 | const std::vector<int>& axes) { |
| 843 | nb::gil_scoped_acquire gil; |
| 844 | |
| 845 | auto new_inputs = nb::cast<nb::tuple>( |
| 846 | tree_unflatten_from_structure(input_structure_, inputs)); |
| 847 | auto args = nb::cast<nb::tuple>(new_inputs[0]); |
| 848 | auto kwargs = nb::cast<nb::dict>(new_inputs[1]); |
| 849 | if (kwargs.size() > 0) { |
| 850 | throw std::invalid_argument( |
| 851 | "[custom vmap] Function should only accept positional arguments"); |
| 852 | } |
| 853 | |
| 854 | int arr_index = 0; |
| 855 | auto new_axes = |
| 856 | nb::cast<nb::tuple>(tree_map(args, [&](nb::handle element) { |
| 857 | int axis = axes[arr_index++]; |
| 858 | if (nb::isinstance<mx::array>(element) && axis >= 0) { |
| 859 | return nb::cast(axis); |
| 860 | } else { |
| 861 | return nb::none(); |
| 862 | } |
| 863 | })); |
| 864 | |
| 865 | nb::object result; |
| 866 | if (args.size() == 1) { |
| 867 | result = vmap_fun_(args[0], new_axes[0]); |
| 868 | } else { |
| 869 | result = vmap_fun_(args, new_axes); |
| 870 | } |
| 871 | |
| 872 | if (!nb::isinstance<nb::tuple>(result)) { |
| 873 | throw std::invalid_argument( |
| 874 | "[custom vmap] Vmap function should return a tuple with 2 items."); |
| 875 | } |
| 876 | nb::tuple result_tuple = nb::cast<nb::tuple>(result); |
| 877 | if (result_tuple.size() != 2) { |
| 878 | throw std::invalid_argument( |
| 879 | "[custom vmap] Vmap function should return a tuple with 2 items."); |
| 880 | } |
| 881 | |
| 882 | std::vector<mx::array> outputs; |
| 883 | std::vector<int> output_axes; |
| 884 | tree_visit({result_tuple[0], result_tuple[1]}, [&](auto objects) { |
| 885 | if (nb::isinstance<mx::array>(objects[0])) { |
| 886 | outputs.push_back(nb::cast<mx::array>(objects[0])); |
| 887 | output_axes.push_back( |
| 888 | objects[1].is_none() ? -1 : nb::cast<int>(objects[1])); |
| 889 | } |
| 890 | }); |
| 891 | |
| 892 | return {outputs, output_axes}; |
| 893 | } |
| 894 | }; |
| 895 | |
| 896 | nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) { |
nothing calls this directly
no test coverage detected