Add tensor split

src/core/include/tensor.h

@@ -225,11 +225,12 @@ class Tensor {
     Tensor contiguous() const;
     Tensor as_strided(std::vector<int64_t> sizes, std::vector<int64_t> strides, int64_t storage_offset = 0) const;
     Tensor permute(const std::vector<int64_t> dims) const;
-    Tensor slice(int64_t dim, std::optional<int64_t> start, std::optional<int64_t> end, int64_t step) const;
+    Tensor slice(int64_t dim, std::optional<int64_t> start, std::optional<int64_t> end, int64_t step = 1) const;
     Tensor select(int64_t dim, int64_t index) const;
     Tensor narrow(int64_t dim, int64_t start, int64_t length) const;
     Tensor view(std::vector<int64_t> sizes) const;
     bool can_use_32bit_indexing() const;
+    std::vector<Tensor> split(std::vector<int64_t> indices, int64_t dim) const;
 
     Tensor _half() const;
     Tensor _float() const;

src/core/include/tensor_shape.h

@@ -7,5 +7,6 @@
 namespace gpu {
 
 Tensor concat(const std::vector<Tensor> tensors, int64_t dim);
+std::vector<Tensor> tensor_split(const Tensor &self, std::vector<int64_t> indices, int64_t dim);
 
-}
+} // namespace gpu

src/core/tensor.cpp

@@ -3,6 +3,7 @@
 #include <vector>
 
 #include "tensor.h"
+#include "tensor_shape.h"
 #include "exception.h"
 #include "data_ptr.h"
 #include "intrusive_ptr.h"
@@ -282,6 +283,10 @@ bool Tensor::can_use_32bit_indexing() const {
     return true;
 }
 
+std::vector<Tensor> Tensor::split(std::vector<int64_t> indices, int64_t dim) const {
+    return gpu::tensor_split(*this, indices, dim);
+}
+
 Tensor Tensor::_half() const {
     return gpu::convert(*this, ScalarType::Half);
 }

src/core/tensor_shape.cpp

@@ -69,4 +69,23 @@ Tensor concat(const std::vector<Tensor> tensors, int64_t dim) {
     return result;
 }
 
+std::vector<Tensor> tensor_split(const Tensor &self, std::vector<int64_t> indices, int64_t dim) {
+    CHECK_FAIL(
+        self.dim() > 0,
+        "tensor_split expected at least a 1-dimensional tensor, but got a tensor with ",
+        self.dim(),
+        " dims");
+    int64_t dim_ = maybe_wrap_dim(dim, self.dim());
+    int64_t num_indices = indices.size();
+    std::vector<Tensor> splits(num_indices);
+    int64_t start_idx = 0;
+    for (int split_idx = 0; split_idx < num_indices; ++split_idx) {
+        auto end_idx = start_idx + indices[split_idx];
+        splits[split_idx] = self.slice(dim_, start_idx, end_idx);
+        start_idx = end_idx;
+    }
+    CHECK_FAIL(start_idx == self.shape(dim));
+    return splits;
+}
+
 } // namespace gpu

src/register.cpp

@@ -134,6 +134,7 @@ PYBIND11_MODULE(kfunca, m) {
             }
             return self.view(dims);
         })
+        .def("split", &Tensor::split)
         .def("sort", &Tensor::sort)
         .def("topk", &Tensor::topk)
         .def("__getitem__", [](Tensor &self, py::object key) {

test/test_tensor.py

@@ -252,6 +252,16 @@ def test_cat(self):
         arr_t = torch.cat([arr1_t, arr2_t, arr3_t], 1)
         arr_gpu = kfunca.cat([arr1_gpu, arr2_gpu, arr3_gpu], 1)
         assert_allclose(arr_t, arr_gpu)
+
+    def test_split(self):
+        arr = np.random.uniform(-10000, 10000, size=(5,25,23)).astype(np.float32)
+        arr_t = torch.from_numpy(arr)
+        arr_gpu = kfunca.from_numpy(arr, 0)
+        arr_t1, arr_t2, arr_t3 = arr_t.split([11,13,1], 1)
+        arr_gpu1, arr_gpu2, arr_gpu3 = arr_gpu.split([11,13,1], 1)
+        assert_allclose(arr_t1, arr_gpu1)
+        assert_allclose(arr_t2, arr_gpu2)
+        assert_allclose(arr_t3, arr_gpu3)
 
 
 if __name__ == '__main__':