| 721 | }; |
| 722 | |
| 723 | struct InnerVJPFunction { |
| 724 | nb::callable vjp_fun_; |
| 725 | nb::object input_structure_; |
| 726 | std::shared_ptr<nb::object> output_structure_; |
| 727 | |
| 728 | InnerVJPFunction( |
| 729 | nb::callable vjp_fun, |
| 730 | nb::object input_structure, |
| 731 | std::shared_ptr<nb::object> output_structure) |
| 732 | : vjp_fun_(std::move(vjp_fun)), |
| 733 | input_structure_(std::move(input_structure)), |
| 734 | output_structure_(std::move(output_structure)) {} |
| 735 | ~InnerVJPFunction() { |
| 736 | nb::gil_scoped_acquire gil; |
| 737 | |
| 738 | vjp_fun_.reset(); |
| 739 | input_structure_.reset(); |
| 740 | if (output_structure_.use_count() == 1) { |
| 741 | output_structure_->reset(); |
| 742 | } |
| 743 | } |
| 744 | |
| 745 | std::vector<mx::array> operator()( |
| 746 | const std::vector<mx::array>& primals, |
| 747 | const std::vector<mx::array>& cotangents, |
| 748 | const std::vector<mx::array>& outputs) { |
| 749 | nb::gil_scoped_acquire gil; |
| 750 | |
| 751 | auto new_inputs = nb::cast<nb::tuple>( |
| 752 | tree_unflatten_from_structure(input_structure_, primals)); |
| 753 | auto args = nb::cast<nb::tuple>(new_inputs[0]); |
| 754 | auto new_cotangents = |
| 755 | tree_unflatten_from_structure(*output_structure_, cotangents); |
| 756 | auto new_outputs = |
| 757 | tree_unflatten_from_structure(*output_structure_, outputs); |
| 758 | |
| 759 | if (args.size() == 1) { |
| 760 | return tree_flatten( |
| 761 | vjp_fun_(args[0], new_cotangents, new_outputs, **new_inputs[1]), |
| 762 | false); |
| 763 | } else { |
| 764 | return tree_flatten( |
| 765 | vjp_fun_(args, new_cotangents, new_outputs, **new_inputs[1]), |
| 766 | false); |
| 767 | } |
| 768 | } |
| 769 | }; |
| 770 | |
| 771 | struct InnerJVPFunction { |
| 772 | nb::callable jvp_fun_; |