VTEN-19-Add-broadcast-to-tensor-assignment (#27)

Author: shengtsui

docs/source/api/core/proxy_slice.rst

@@ -10,6 +10,9 @@ vt::TensorSliceProxy
 .. doxygenstruct:: vt::NewAxisT
    :members:
 
+.. doxygenenum:: vt::SliceType
+   :members:
+
 .. doxygenclass:: vt::Slice
    :members:
 

docs/source/index.rst

@@ -55,6 +55,8 @@ Comparison between NumPy and VTensor
      - vt::expand_dims_lhs<T, N, 2>(tensor)
    * - arr[..., None, None]
      - vt::expand_dims_rhs<T, N, 2>(tensor)
+   * - arr[:, :]
+     - tensor(vt::Slice::all(), vt::Slice::all())
 
 .. list-table:: Broadcasting
    :header-rows: 1

lib/core/slice.hpp

@@ -1,22 +1,30 @@
 #pragma once
 
 #include "lib/core/assertions.hpp"
+#include "lib/core/broadcast.hpp"
 #include "lib/core/tensor.hpp"
 
 namespace vt {
 
+/**
+ * @brief Slice type.
+ */
+enum class SliceType { Normal, All };
+
 /**
  * @brief Slice class to represent a slice of a tensor.
  * Usage:
  * Slice(0) -> Slice from 0 to 1 with step 1
  * Slice(0, 10) -> Slice from 0 to 10 with step 1
  * Slice(0, 10, 2) -> Slice from 0 to 10 with step 2
+ * Slice::all() -> Slice to represent the full range
  */
 class Slice {
    public:
     size_t start;
     size_t end;
     size_t step;
+    SliceType type = SliceType::Normal;
 
     // Default constructor
     Slice() = default;
@@ -44,6 +52,20 @@ class Slice {
      * @param step: The step of the slice.
      */
     Slice(size_t start, size_t end, size_t step) : start(start), end(end), step(step) {}
+
+    /**
+     * @brief Construct a new Slice object
+     *
+     * @param type: The type of the slice.
+     */
+    Slice(SliceType type) : type(type) {}
+
+    /**
+     * @brief Return a slice object to represent the full range.
+     *
+     * @return Slice: The slice object.
+     */
+    static Slice all() { return Slice(SliceType::All); }
 };
 
 /**
@@ -99,8 +121,8 @@ class TensorSliceProxy : public Tensor<T, N> {
      */
     TensorSliceProxy& operator=(const Tensor<T, N>& other) {
         assert_same_order_between_two_tensors(this->order(), other.order());
-        assert(this->shape() == other.shape());
-        thrust::copy(other.begin(), other.end(), this->begin());
+        auto [lhs, rhs] = broadcast(*this, other);
+        thrust::copy(rhs.begin(), rhs.end(), lhs.begin());
         return *this;
     }
 

lib/core/tensor.hpp

@@ -417,11 +417,16 @@ class Tensor {
         Shape<N> new_shape{};
         size_t offset = 0;
         for (size_t i = 0; i < N; i++) {
-            assert(slices[i].end <= _shape[i]);       // Check if the end of the slice is within the tensor dimension.
-            assert(slices[i].start < slices[i].end);  // Check if the start of the slice is less than the end.
-            new_strides[i] = slices[i].step * _strides[i];
-            new_shape[i] = (slices[i].end - slices[i].start + slices[i].step - 1) / slices[i].step;
-            offset += slices[i].start * _strides[i];
+            if (slices[i].type == SliceType::All) {
+                new_strides[i] = _strides[i];
+                new_shape[i] = _shape[i];
+            } else {
+                assert(slices[i].end <= _shape[i]);       // Check if the end of the slice is within the tensor dimension.
+                assert(slices[i].start < slices[i].end);  // Check if the start of the slice is less than the end.
+                new_strides[i] = slices[i].step * _strides[i];
+                new_shape[i] = (slices[i].end - slices[i].start + slices[i].step - 1) / slices[i].step;
+                offset += slices[i].start * _strides[i];
+            }
         }
         size_t new_start = _start + offset;
         return Tensor<T, N>(_data, new_shape, new_strides, new_start, _order, false);
@@ -443,7 +448,13 @@ class Tensor {
      *
      * @return T*: The raw pointer of the tensor.
      */
-    T* raw_ptr() const { return thrust::raw_pointer_cast(_data->data()); }
+    T* raw_ptr() const {
+        if (_data == nullptr) {
+            return nullptr;
+        } else {
+            return thrust::raw_pointer_cast(_data->data());
+        }
+    }
 
     /**
      * @brief Return the contiguous flag of the tensor.

lib/core/tests/test_slice.cc

@@ -135,4 +135,18 @@ TEST(TensorCondProxyFromConstantF, BasicAssertions) {
     auto tensor = vt::arange(24, vt::Order::F)({0, 24, 2}).reshape(3, 4);
     tensor[tensor > 12.0f] = 1.0f;
     EXPECT_EQ(vt::asvector(tensor), (std::vector<float>{0, 2, 4, 6, 8, 10, 12, 1, 1, 1, 1, 1}));
-}
+}
+
+TEST(BroadcastSliceAssignmentC, BasicAssertions) {
+    auto tensor = vt::arange(12).reshape(4, 3);
+    tensor(vt::Slice::all(), vt::Slice(0, 2, 1)) = tensor(vt::Slice::all(), vt::Slice(2, 3, 1));
+    EXPECT_EQ(vt::asvector(tensor), (std::vector<float>{2, 2, 2, 5, 5, 5, 8, 8, 8, 11, 11, 11}));
+    EXPECT_EQ(tensor.contiguous(), true);
+}
+
+TEST(BroadcastSliceAssignmentF, BasicAssertions) {
+    auto tensor = vt::arange(12, vt::Order::F).reshape(4, 3);
+    tensor(vt::Slice::all(), vt::Slice(0, 2, 1)) = tensor(vt::Slice::all(), vt::Slice(2, 3, 1));
+    EXPECT_EQ(vt::asvector(tensor), (std::vector<float>{8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11}));
+    EXPECT_EQ(tensor.contiguous(), true);
+}