Function TEST

test/cpp/compat/ATen_split_test.cc:34–49

Test for tensor_split with sections

// Test for tensor_split with sections
TEST(TensorSplitTest, TensorSplitWithSections) {
  // Create a test tensor [0, 1, 2, ..., 8] (9 elements, evenly divisible by 3)
  auto tensor = at::arange(9, at::TensorOptions().dtype(at::kFloat));

  // Split into 3 sections along dim 0
  auto splits = tensor.tensor_split(3, 0);

  EXPECT_EQ(splits.size(), 3);
  EXPECT_EQ(splits[0].numel(), 3);  // [0, 1, 2]
  EXPECT_EQ(splits[1].numel(), 3);  // [3, 4, 5]
  EXPECT_EQ(splits[2].numel(), 3);  // [6, 7, 8]

  // Verify first split values
  EXPECT_FLOAT_EQ(splits[0][0].item<float>(), 0.0f);
  EXPECT_FLOAT_EQ(splits[0][2].item<float>(), 2.0f);
}

// Test for tensor_split with indices (PyTorch semantics: indices are positions)
TEST(TensorSplitTest, TensorSplitWithIndices) {