| 687 | } |
| 688 | |
| 689 | struct InnerFunction { |
| 690 | nb::callable fun_; |
| 691 | nb::object input_structure_; |
| 692 | std::shared_ptr<nb::object> output_structure_; |
| 693 | |
| 694 | InnerFunction( |
| 695 | nb::callable fun, |
| 696 | nb::object input_structure, |
| 697 | std::shared_ptr<nb::object> output_structure) |
| 698 | : fun_(std::move(fun)), |
| 699 | input_structure_(std::move(input_structure)), |
| 700 | output_structure_(std::move(output_structure)) {} |
| 701 | ~InnerFunction() { |
| 702 | nb::gil_scoped_acquire gil; |
| 703 | |
| 704 | fun_.reset(); |
| 705 | input_structure_.reset(); |
| 706 | if (output_structure_.use_count() == 1) { |
| 707 | output_structure_->reset(); |
| 708 | } |
| 709 | } |
| 710 | |
| 711 | std::vector<mx::array> operator()(const std::vector<mx::array>& inputs) { |
| 712 | nb::gil_scoped_acquire gil; |
| 713 | |
| 714 | auto new_inputs = nb::cast<nb::tuple>( |
| 715 | tree_unflatten_from_structure(input_structure_, inputs)); |
| 716 | std::vector<mx::array> outputs; |
| 717 | std::tie(outputs, *output_structure_) = |
| 718 | tree_flatten_with_structure(fun_(*new_inputs[0], **new_inputs[1])); |
| 719 | return outputs; |
| 720 | } |
| 721 | }; |
| 722 | |
| 723 | struct InnerVJPFunction { |
| 724 | nb::callable vjp_fun_; |