| 824 | }; |
| 825 | |
| 826 | struct InnerVmapFunction { |
| 827 | nb::callable vmap_fun_; |
| 828 | nb::object input_structure_; |
| 829 | |
| 830 | InnerVmapFunction(nb::callable vmap_fun, nb::object input_structure) |
| 831 | : vmap_fun_(std::move(vmap_fun)), |
| 832 | input_structure_(std::move(input_structure)) {} |
| 833 | ~InnerVmapFunction() { |
| 834 | nb::gil_scoped_acquire gil; |
| 835 | |
| 836 | vmap_fun_.reset(); |
| 837 | input_structure_.reset(); |
| 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; |