From 22a05d504238b99dcebb359401379aaab33ac9de Mon Sep 17 00:00:00 2001
From: Jingyue Wu <wujingyue@gmail.com>
Date: Tue, 26 Nov 2024 10:44:41 -0800
Subject: [PATCH] Add a repro for #3479

---
 tests/cpp/test_allocation_domain.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/cpp/test_allocation_domain.cpp b/tests/cpp/test_allocation_domain.cpp
index a4c6829d122..e7ee94c64ca 100644
--- a/tests/cpp/test_allocation_domain.cpp
+++ b/tests/cpp/test_allocation_domain.cpp
@@ -1426,8 +1426,9 @@ TEST_F(AllocationDomainTest, InputAllocationIsSplit_Concrete) {
   fusion->addInput(in);
   fusion->addOutput(out);
 
-  in->split(0, 2);
-  in->setAllocationDomain(in->getLoopDomain(), true);
+  auto [outer, inner] = IterDomain::split(
+      in->axis(0), IrBuilder::create<Val>(2, DataType::Index), true);
+  in->setAllocationDomain({outer, inner}, true);
 
   FusionExecutorCache executor_cache(std::move(fusion));
   auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA);