add untested tests

tenstorrent · Mar 6, 2025 · 4e5f8a9 · 4e5f8a9
1 parent e2a46b1
commit 4e5f8a9
Showing 1 changed file with 42 additions and 0 deletions.
diff --git a/tests/ttnn/unit_tests/gtests/tensor/test_distributed_tensor.cpp b/tests/ttnn/unit_tests/gtests/tensor/test_distributed_tensor.cpp
@@ -22,6 +22,48 @@ TensorSpec get_tensor_spec(const ttnn::Shape& shape, DataType dtype) {
     return TensorSpec(shape, TensorLayout(dtype, Layout::ROW_MAJOR, MemoryConfig{}));
 }
 
+TEST_F(TensorDistributionTest, DeviceAggregate) {
+    const int num_devices = mesh_device_->num_devices();
+    std::vector<std::vector<float>> test_data(num_devices);
+    for (int i = 0; i < num_devices; i++) {
+        test_data[i].insert(test_data[i].end(), {i * 1.F, i * 2.F, i * 3.F});
+    }
+
+    std::vector<Tensor> tensors(num_devices);
+
+    for(int i = 0; i < num_devices; i++) {
+        tensors.push_back(Tensor::from_vector(test_data[i], get_tensor_spec(ttnn::Shape{1, num_devices, 3, 1}, DataType::FLOAT32)), mesh_device_);
+    }
+
+    Tensor aggregated_tensor = aggregate_as_tensor(tensors, AllGatherTensor{});
+    System.out.println(aggregated_tensor)
+    EXPECT_TRUE(aggregated_tensor.storage_type() == StorageType::MULTI_DEVICE);
+
+    std::vector<float> out_vector = aggregated_tensor.to_vector<float>();
+    EXPECT_EQ(out_vector, test_data);
+}
+
+TEST_F(TensorDistributionTest, BorrowedAggregate) {
+    const int num_devices = mesh_device_->num_devices();
+    std::vector<std::vector<float>> test_data(num_devices);
+    for (int i = 0; i < num_devices; i++) {
+        test_data[i].insert(test_data[i].end(), {i * 1.F, i * 2.F, i * 3.F});
+    }
+
+    std::vector<Tensor> tensors(num_devices);
+
+    for(int i = 0; i < num_devices; i++) {
+        tensors.push_back(Tensor()  Tensor::from_vector(test_data[i], get_tensor_spec(ttnn::Shape{1, num_devices, 3, 1}, DataType::FLOAT32)));
+    }
+Storage storage, const ttnn::Shape& shape, DataType dtype, Layout layout, const std::optional<Tile>& tile
+    Tensor aggregated_tensor = aggregate_as_tensor(tensors, AllGatherTensor{});
+    System.out.println(aggregated_tensor)
+    EXPECT_TRUE(aggregated_tensor.storage_type() == StorageType::MULTI_DEVICE_HOST);
+
+    std::vector<float> out_vector = aggregated_tensor.to_vector<float>();
+    EXPECT_EQ(out_vector, test_data);
+}
+
 TEST_F(TensorDistributionTest, DistributeToDevice) {
     Tensor input_tensor = Tensor::from_vector(
         std::vector<float>{42.F, 13.F, -99.F}, get_tensor_spec(ttnn::Shape{1, 1, 1, 3}, DataType::FLOAT32));