diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index 5a17345..ccac42a 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -225,11 +225,12 @@ class Tensor { Tensor contiguous() const; Tensor as_strided(std::vector sizes, std::vector strides, int64_t storage_offset = 0) const; Tensor permute(const std::vector dims) const; - Tensor slice(int64_t dim, std::optional start, std::optional end, int64_t step) const; + Tensor slice(int64_t dim, std::optional start, std::optional 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 sizes) const; bool can_use_32bit_indexing() const; + std::vector split(std::vector indices, int64_t dim) const; Tensor _half() const; Tensor _float() const; diff --git a/src/core/include/tensor_shape.h b/src/core/include/tensor_shape.h index 6cbdffb..5376519 100644 --- a/src/core/include/tensor_shape.h +++ b/src/core/include/tensor_shape.h @@ -7,5 +7,6 @@ namespace gpu { Tensor concat(const std::vector tensors, int64_t dim); +std::vector tensor_split(const Tensor &self, std::vector indices, int64_t dim); -} +} // namespace gpu diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index 763a217..af3451d 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -3,6 +3,7 @@ #include #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::split(std::vector indices, int64_t dim) const { + return gpu::tensor_split(*this, indices, dim); +} + Tensor Tensor::_half() const { return gpu::convert(*this, ScalarType::Half); } diff --git a/src/core/tensor_shape.cpp b/src/core/tensor_shape.cpp index 0b645ed..089be1d 100644 --- a/src/core/tensor_shape.cpp +++ b/src/core/tensor_shape.cpp @@ -69,4 +69,23 @@ Tensor concat(const std::vector tensors, int64_t dim) { return result; } +std::vector tensor_split(const Tensor &self, std::vector 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 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 diff --git a/src/register.cpp b/src/register.cpp index 4a93112..c05adf3 100644 --- a/src/register.cpp +++ b/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) { diff --git a/test/test_tensor.py b/test/test_tensor.py index 45aae5d..4516428 100644 --- a/test/test_tensor.py +++ b/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__':