Add move and swap axis, and vmap for slice, concat, and gather (#158)

* add move and swap axis, and vmap for slice, concat, and gather
2025-12-14 09:07:12 +08:00 · 2023-12-14 12:59:12 -08:00
parent f55908bc48
commit e5851e52b1
10 changed files with 399 additions and 7 deletions
--- a/python/src/array.cpp
+++ b/python/src/array.cpp
@@ -862,6 +862,22 @@ void init_array(py::module_& m) {
          py::kw_only(),
          "stream"_a = none,
          "See :func:`any`.")
+      .def(
+          "moveaxis",
+          &moveaxis,
+          "source"_a,
+          "destination"_a,
+          py::kw_only(),
+          "stream"_a = none,
+          "See :func:`moveaxis`.")
+      .def(
+          "swapaxes",
+          &swapaxes,
+          "axis1"_a,
+          "axis2"_a,
+          py::kw_only(),
+          "stream"_a = none,
+          "See :func:`moveaxis`.")
      .def(
          "transpose",
          [](const array& a, py::args axes, StreamOrDevice s) {