| 21 | using namespace nb::literals; |
| 22 | |
| 23 | std::pair<mx::Args, mx::Kwargs> validate_and_extract_inputs( |
| 24 | const nb::args& args, |
| 25 | const nb::kwargs& kwargs, |
| 26 | const std::string& prefix) { |
| 27 | auto maybe_throw = [&prefix](bool valid) { |
| 28 | if (!valid) { |
| 29 | throw std::invalid_argument( |
| 30 | prefix + |
| 31 | " Inputs can either be a variable " |
| 32 | "number of positional and keyword arrays or a single tuple " |
| 33 | "and/or dictionary of arrays."); |
| 34 | } |
| 35 | }; |
| 36 | mx::Args args_; |
| 37 | mx::Kwargs kwargs_; |
| 38 | if (args.size() == 0) { |
| 39 | // No args so kwargs must be keyword arrays |
| 40 | maybe_throw(nb::try_cast(kwargs, kwargs_)); |
| 41 | } else if (args.size() > 0 && nb::isinstance<mx::array>(args[0])) { |
| 42 | // Args are positional arrays and kwargs are keyword arrays |
| 43 | maybe_throw(nb::try_cast(args, args_)); |
| 44 | maybe_throw(nb::try_cast(kwargs, kwargs_)); |
| 45 | } else if (args.size() == 1) { |
| 46 | // - args[0] can be a tuple or list or arrays or a dict |
| 47 | // with string keys and array values |
| 48 | // - kwargs should be empty |
| 49 | maybe_throw(kwargs.size() == 0); |
| 50 | if (!nb::try_cast(args[0], args_)) { |
| 51 | maybe_throw(nb::try_cast(args[0], kwargs_)); |
| 52 | } |
| 53 | } else if (args.size() == 2) { |
| 54 | // - args[0] can be a tuple or list of arrays |
| 55 | // - args[1] can be a dict of string keys with array values. |
| 56 | // - kwargs should be empty |
| 57 | maybe_throw(kwargs.size() == 0); |
| 58 | maybe_throw(nb::try_cast(args[0], args_)); |
| 59 | maybe_throw(nb::try_cast(args[1], kwargs_)); |
| 60 | } else { |
| 61 | maybe_throw(false); |
| 62 | } |
| 63 | return {args_, kwargs_}; |
| 64 | } |
| 65 | |
| 66 | int py_function_exporter_tp_traverse( |
| 67 | PyObject* self, |