| 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 { |
nothing calls this directly
no test coverage detected