| 769 | }; |
| 770 | |
| 771 | struct InnerJVPFunction { |
| 772 | nb::callable jvp_fun_; |
| 773 | nb::object input_structure_; |
| 774 | |
| 775 | InnerJVPFunction(nb::callable jvp_fun, nb::object input_structure) |
| 776 | : jvp_fun_(std::move(jvp_fun)), |
| 777 | input_structure_(std::move(input_structure)) {} |
| 778 | ~InnerJVPFunction() { |
| 779 | nb::gil_scoped_acquire gil; |
| 780 | |
| 781 | jvp_fun_.reset(); |
| 782 | input_structure_.reset(); |
| 783 | } |
| 784 | |
| 785 | std::vector<mx::array> operator()( |
| 786 | const std::vector<mx::array>& primals, |
| 787 | const std::vector<mx::array>& tangents, |
| 788 | const std::vector<int>& argnums) { |
| 789 | nb::gil_scoped_acquire gil; |
| 790 | |
| 791 | auto new_inputs = nb::cast<nb::tuple>( |
| 792 | tree_unflatten_from_structure(input_structure_, primals)); |
| 793 | auto args = nb::cast<nb::tuple>(new_inputs[0]); |
| 794 | auto kwargs = nb::cast<nb::dict>(new_inputs[1]); |
| 795 | if (kwargs.size() > 0) { |
| 796 | throw std::invalid_argument( |
| 797 | "[custom jvp] Function should only accept positional arguments"); |
| 798 | } |
| 799 | |
| 800 | // Make a new pytree which has tangents or None when a tangent is not |
| 801 | // available. |
| 802 | std::vector<bool> have_tangents(primals.size(), false); |
| 803 | for (auto arg : argnums) { |
| 804 | have_tangents[arg] = true; |
| 805 | } |
| 806 | int array_index = 0; |
| 807 | int tangent_index = 0; |
| 808 | auto new_tangents = |
| 809 | nb::cast<nb::tuple>(tree_map(args, [&](nb::handle element) { |
| 810 | if (nb::isinstance<mx::array>(element) && |
| 811 | have_tangents[array_index++]) { |
| 812 | return nb::cast(tangents[tangent_index++]); |
| 813 | } else { |
| 814 | return nb::none(); |
| 815 | } |
| 816 | })); |
| 817 | |
| 818 | if (args.size() == 1) { |
| 819 | return tree_flatten(jvp_fun_(args[0], new_tangents[0]), false); |
| 820 | } else { |
| 821 | return tree_flatten(jvp_fun_(args, new_tangents), false); |
| 822 | } |
| 823 | } |
| 824 | }; |
| 825 | |
| 826 | struct InnerVmapFunction { |
| 827 | nb::callable vmap_fun_; |