| 307 | } |
| 308 | |
| 309 | auto py_vmap( |
| 310 | const nb::callable& fun, |
| 311 | const nb::object& in_axes, |
| 312 | const nb::object& out_axes) { |
| 313 | return [fun, in_axes, out_axes](const nb::args& args) { |
| 314 | auto axes_to_flat_tree = [](const nb::object& tree, |
| 315 | const nb::object& axes, |
| 316 | bool output_axes) { |
| 317 | std::vector<int> flat_axes; |
| 318 | bool encountered_tuple = false; |
| 319 | tree_visit( |
| 320 | {tree, axes}, |
| 321 | [&flat_axes, &encountered_tuple, output_axes]( |
| 322 | const std::vector<nb::object>& inputs) { |
| 323 | if (nb::isinstance<mx::array>(inputs[0])) { |
| 324 | if (inputs[1].is_none()) { |
| 325 | flat_axes.push_back(-1); |
| 326 | } else if (nb::isinstance<nb::int_>(inputs[1])) { |
| 327 | int axis = nb::cast<int>(nb::cast<nb::int_>(inputs[1])); |
| 328 | const mx::array& x = nb::cast<mx::array>(inputs[0]); |
| 329 | if (axis < 0) { |
| 330 | axis += x.ndim() + output_axes; |
| 331 | } |
| 332 | if (axis < 0 || axis >= (x.ndim() + output_axes)) { |
| 333 | std::ostringstream msg; |
| 334 | msg << "[vmap] Invalid" << (output_axes ? " output " : " ") |
| 335 | << "vectorization axis " << axis |
| 336 | << " for array with shape " << x.shape(); |
| 337 | throw std::invalid_argument(msg.str()); |
| 338 | } |
| 339 | flat_axes.push_back(axis); |
| 340 | } else if (nb::isinstance<nb::tuple>(inputs[1])) { |
| 341 | encountered_tuple = true; |
| 342 | auto l = nb::cast<nb::tuple>(inputs[1]); |
| 343 | if (l.size() == 1 && nb::isinstance<nb::int_>(l[0])) { |
| 344 | int axis = nb::cast<int>(nb::cast<nb::int_>(l[0])); |
| 345 | const mx::array& x = nb::cast<mx::array>(inputs[0]); |
| 346 | if (axis < 0) { |
| 347 | axis += x.ndim() + output_axes; |
| 348 | } |
| 349 | if (axis < 0 || axis >= (x.ndim() + output_axes)) { |
| 350 | std::ostringstream msg; |
| 351 | msg << "[vmap] Invalid" << (output_axes ? " output " : " ") |
| 352 | << "vectorization axis " << axis |
| 353 | << " for array with shape " << x.shape(); |
| 354 | throw std::invalid_argument(msg.str()); |
| 355 | } |
| 356 | flat_axes.push_back(axis); |
| 357 | } else if (l.size() == 1 && l[0].is_none()) { |
| 358 | flat_axes.push_back(-1); |
| 359 | } else { |
| 360 | throw std::invalid_argument( |
| 361 | "[vmap] axis must be int or None."); |
| 362 | } |
| 363 | } else { |
| 364 | throw std::invalid_argument("[vmap] axis must be int or None."); |
| 365 | } |
| 366 | } else { |
no test coverage detected