add squeeze and unsqueeze test

Author: youge325

test/TensorTest.cpp

@@ -170,5 +170,108 @@ TEST_F(TensorTest, Transpose) {
   EXPECT_EQ(transposed.sizes()[2], 2);
 }
 
+// 测试 squeeze
+TEST_F(TensorTest, Squeeze) {
+  // 创建一个包含大小为1的维度的tensor: shape = {2, 1, 3, 1, 4}
+  at::Tensor tensor_with_ones = at::ones({2, 1, 3, 1, 4}, at::kFloat);
+
+  // 移除所有大小为1的维度
+  at::Tensor squeezed = tensor_with_ones.squeeze();
+  EXPECT_EQ(squeezed.dim(), 3);
+  EXPECT_EQ(squeezed.sizes()[0], 2);
+  EXPECT_EQ(squeezed.sizes()[1], 3);
+  EXPECT_EQ(squeezed.sizes()[2], 4);
+  EXPECT_EQ(squeezed.numel(), 24);
+
+  // 移除指定维度（维度1，大小为1）
+  at::Tensor squeezed_dim1 = tensor_with_ones.squeeze(1);
+  EXPECT_EQ(squeezed_dim1.dim(), 4);
+  EXPECT_EQ(squeezed_dim1.sizes()[0], 2);
+  EXPECT_EQ(squeezed_dim1.sizes()[1], 3);
+  EXPECT_EQ(squeezed_dim1.sizes()[2], 1);
+  EXPECT_EQ(squeezed_dim1.sizes()[3], 4);
+}
+
+// 测试 squeeze_ (原位操作)
+TEST_F(TensorTest, SqueezeInplace) {
+  // 创建一个包含大小为1的维度的tensor: shape = {2, 1, 3, 1, 4}
+  at::Tensor tensor_with_ones = at::ones({2, 1, 3, 1, 4}, at::kFloat);
+
+  // 记录原始数据指针
+  void* original_ptr = tensor_with_ones.data_ptr();
+
+  // 原位移除所有大小为1的维度
+  tensor_with_ones.squeeze_();
+  EXPECT_EQ(tensor_with_ones.dim(), 3);
+  EXPECT_EQ(tensor_with_ones.sizes()[0], 2);
+  EXPECT_EQ(tensor_with_ones.sizes()[1], 3);
+  EXPECT_EQ(tensor_with_ones.sizes()[2], 4);
+  EXPECT_EQ(tensor_with_ones.numel(), 24);
+
+  // 验证是原位操作（数据指针未改变）
+  EXPECT_EQ(tensor_with_ones.data_ptr(), original_ptr);
+
+  // 测试原位移除指定维度
+  at::Tensor tensor_with_ones2 = at::ones({2, 1, 3, 1, 4}, at::kFloat);
+  tensor_with_ones2.squeeze_(1);
+  EXPECT_EQ(tensor_with_ones2.dim(), 4);
+  EXPECT_EQ(tensor_with_ones2.sizes()[1], 3);
+}
+
+// 测试 unsqueeze
+TEST_F(TensorTest, Unsqueeze) {
+  // 在维度0之前添加一个大小为1的维度
+  at::Tensor unsqueezed0 = tensor.unsqueeze(0);
+  EXPECT_EQ(unsqueezed0.dim(), 4);
+  EXPECT_EQ(unsqueezed0.sizes()[0], 1);
+  EXPECT_EQ(unsqueezed0.sizes()[1], 2);
+  EXPECT_EQ(unsqueezed0.sizes()[2], 3);
+  EXPECT_EQ(unsqueezed0.sizes()[3], 4);
+  EXPECT_EQ(unsqueezed0.numel(), 24);
+
+  // 在维度2之前添加一个大小为1的维度
+  at::Tensor unsqueezed2 = tensor.unsqueeze(2);
+  EXPECT_EQ(unsqueezed2.dim(), 4);
+  EXPECT_EQ(unsqueezed2.sizes()[0], 2);
+  EXPECT_EQ(unsqueezed2.sizes()[1], 3);
+  EXPECT_EQ(unsqueezed2.sizes()[2], 1);
+  EXPECT_EQ(unsqueezed2.sizes()[3], 4);
+
+  // 在最后添加一个大小为1的维度（使用负索引-1）
+  at::Tensor unsqueezed_last = tensor.unsqueeze(-1);
+  EXPECT_EQ(unsqueezed_last.dim(), 4);
+  EXPECT_EQ(unsqueezed_last.sizes()[0], 2);
+  EXPECT_EQ(unsqueezed_last.sizes()[1], 3);
+  EXPECT_EQ(unsqueezed_last.sizes()[2], 4);
+  EXPECT_EQ(unsqueezed_last.sizes()[3], 1);
+}
+
+// 测试 unsqueeze_ (原位操作)
+TEST_F(TensorTest, UnsqueezeInplace) {
+  // 创建一个新的tensor用于原位操作
+  at::Tensor test_tensor = at::ones({2, 3, 4}, at::kFloat);
+
+  // 记录原始数据指针
+  void* original_ptr = test_tensor.data_ptr();
+
+  // 原位在维度0之前添加一个大小为1的维度
+  test_tensor.unsqueeze_(0);
+  EXPECT_EQ(test_tensor.dim(), 4);
+  EXPECT_EQ(test_tensor.sizes()[0], 1);
+  EXPECT_EQ(test_tensor.sizes()[1], 2);
+  EXPECT_EQ(test_tensor.sizes()[2], 3);
+  EXPECT_EQ(test_tensor.sizes()[3], 4);
+  EXPECT_EQ(test_tensor.numel(), 24);
+
+  // 验证是原位操作（数据指针未改变）
+  EXPECT_EQ(test_tensor.data_ptr(), original_ptr);
+
+  // 测试使用负索引的原位操作
+  at::Tensor test_tensor2 = at::ones({2, 3, 4}, at::kFloat);
+  test_tensor2.unsqueeze_(-1);
+  EXPECT_EQ(test_tensor2.dim(), 4);
+  EXPECT_EQ(test_tensor2.sizes()[3], 1);
+}
+
 }  // namespace test
 }  // namespace at