| 132 | } |
| 133 | |
| 134 | void init_export(nb::module_& m) { |
| 135 | m.def( |
| 136 | "export_function", |
| 137 | [](nb::object& file_or_callback, |
| 138 | const nb::callable& fun, |
| 139 | const nb::args& args, |
| 140 | bool shapeless, |
| 141 | const nb::kwargs& kwargs) { |
| 142 | auto [args_, kwargs_] = |
| 143 | validate_and_extract_inputs(args, kwargs, "[export_function]"); |
| 144 | if (nb::isinstance<nb::str>(file_or_callback)) { |
| 145 | mx::export_function( |
| 146 | nb::cast<std::string>(file_or_callback), |
| 147 | wrap_export_function(fun), |
| 148 | args_, |
| 149 | kwargs_, |
| 150 | shapeless); |
| 151 | } else { |
| 152 | auto callback = nb::cast<nb::callable>(file_or_callback); |
| 153 | auto wrapped_callback = |
| 154 | [callback](const mx::ExportCallbackInput& input) { |
| 155 | return callback(input); |
| 156 | }; |
| 157 | mx::export_function( |
| 158 | callback, wrap_export_function(fun), args_, kwargs_, shapeless); |
| 159 | } |
| 160 | }, |
| 161 | nb::arg(), |
| 162 | "fun"_a, |
| 163 | "args"_a, |
| 164 | nb::kw_only(), |
| 165 | "shapeless"_a = false, |
| 166 | "kwargs"_a, |
| 167 | nb::sig( |
| 168 | "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, **kwargs) -> None"), |
| 169 | R"pbdoc( |
| 170 | Export an MLX function. |
| 171 | |
| 172 | Example input arrays must be provided to export a function. The example |
| 173 | inputs can be variable ``*args`` and ``**kwargs`` or a tuple of arrays |
| 174 | and/or dictionary of string keys with array values. |
| 175 | |
| 176 | .. warning:: |
| 177 | |
| 178 | This is part of an experimental API which is likely to |
| 179 | change in future versions of MLX. Functions exported with older |
| 180 | versions of MLX may not be compatible with future versions. |
| 181 | |
| 182 | Args: |
| 183 | file_or_callback (str or Callable): Either a file path to export |
| 184 | the function to or a callback. |
| 185 | fun (Callable): A function which takes as input zero or more |
| 186 | :class:`array` and returns one or more :class:`array`. |
| 187 | *args (array): Example array inputs to the function. |
| 188 | shapeless (bool, optional): Whether or not the function allows |
| 189 | inputs with variable shapes. Default: ``False``. |
| 190 | **kwargs (array): Additional example keyword array inputs to the |
| 191 | function. |
no test coverage detected