| 113 | {0, 0}}; |
| 114 | |
| 115 | auto wrap_export_function(nb::callable fun) { |
| 116 | return |
| 117 | [fun = std::move(fun)](const mx::Args& args_, const mx::Kwargs& kwargs_) { |
| 118 | auto kwargs = nb::dict(); |
| 119 | kwargs.update(nb::cast(kwargs_)); |
| 120 | auto args = nb::tuple(nb::cast(args_)); |
| 121 | auto outputs = fun(*args, **kwargs); |
| 122 | std::vector<mx::array> outputs_; |
| 123 | if (nb::isinstance<mx::array>(outputs)) { |
| 124 | outputs_.push_back(nb::cast<mx::array>(outputs)); |
| 125 | } else if (!nb::try_cast(outputs, outputs_)) { |
| 126 | throw std::invalid_argument( |
| 127 | "[export_function] Outputs can be either a single array " |
| 128 | "a tuple or list of arrays."); |
| 129 | } |
| 130 | return outputs_; |
| 131 | }; |
| 132 | } |
| 133 | |
| 134 | void init_export(nb::module_& m) { |
| 135 | m.def( |
no test coverage detected