removing out-dated assert on python API (#1724)

removing out-dated asserts in python API `define_vector`;
adding a tests verifying the behavior

csrc/python_frontend/python_bindings.cpp

@@ -55,9 +55,6 @@ Vector define_vector_fn(
   std::vector<Scalar> args;
   size_t idx = 0;
   for (const auto& item : values) {
-    NVF_CHECK(
-        idx < 8,
-        "The specified vector size exceeds the max tensor size for nvfuser.");
     if (py::isinstance<py::int_>(item)) {
       auto int_value = py::cast<int64_t>(item);
       NVF_CHECK(
@@ -866,9 +863,6 @@ void initNvFuserPythonBindings(PyObject* module) {
   fusion_def.def(
       "define_vector",
       [](FusionDefinition& self, size_t size) -> Vector {
-        NVF_CHECK(
-            size < 8,
-            "The specified vector size exceeds the max tensor size for nvfuser.");
         std::vector<Scalar> args;
         args.reserve(size);
         for (size_t i = 0; i < size; ++i) {

python_tests/pytest_input_generators.py

@@ -26,7 +26,6 @@
 MINIMUM_SYMBOLIC_SIZE = -1
 INT64_MAX = 2**63 - 1
 MAX_TENSOR_DIMS = 8
-MAX_VECTOR_SIZE = 8
 
 
 # Determine if a number is with desired Domain [low, high)
@@ -468,42 +467,11 @@ def define_vector_constant_error_generator(
         "The value -2 at index 0 was neither symbolic(-1), zero_element(0), broadcast(1), or static(>1)",
     )
 
-    check_max_vector_size = ErrorSample(
-        {
-            "values": [-1 for _ in range(MAX_VECTOR_SIZE + 1)],
-        },
-        "The specified vector size exceeds the max tensor size for nvfuser.",
-    )
-
     error_cases = [
         # FIXME: The above_size_range case gives a non-sensical error message.
         # "Unable to cast Python instance to C++ type (#define PYBIND11_DETAILED_ER"
         # check_above_size_range,
         check_below_size_range,
-        check_max_vector_size,
-    ]
-
-    for es in error_cases:
-        yield SampleInput(**es.kwargs), es.ex_type, es.ex_str
-
-
-def define_vector_input_error_generator(
-    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
-):
-    """
-    "define_vector",
-    [](FusionDefinition& self, size_t size) -> Vector {
-    """
-
-    check_max_vector_size = ErrorSample(
-        {
-            "size": (MAX_VECTOR_SIZE + 1),
-        },
-        "The specified vector size exceeds the max tensor size for nvfuser.",
-    )
-
-    error_cases = [
-        check_max_vector_size,
     ]
 
     for es in error_cases:

python_tests/pytest_opinfos.py

@@ -22,7 +22,6 @@
     define_tensor_generator,
     define_tensor_error_generator,
     define_vector_constant_error_generator,
-    define_vector_input_error_generator,
     elementwise_binary_generator,
     _elementwise_binary_torch,
     elementwise_unary_generator,
@@ -90,15 +89,6 @@
 )
 fusion_input_ops.append(define_vector_constant_opinfo)
 
-define_vector_input_opinfo = OpInfo(
-    lambda fd: fd.define_vector,
-    "define_vector_input",
-    sample_input_generator=None,
-    error_input_generator=define_vector_input_error_generator,
-    fd_error_input_fn=api_test_fd_fn,
-)
-fusion_input_ops.append(define_vector_input_opinfo)
-
 """ End Fusion Input Operations """
 
 """ Start Unary-Float Operations """

python_tests/test_python_frontend.py

@@ -519,6 +519,26 @@ def nvfuser_fusion(
 
         self.assertEqual(eager_out, nvf_out[0])
 
+    def test_tensor_ndim(self):
+        shape = [2 for i in range(12)]
+        new_shape = shape[:9]
+        new_shape.append(8)
+
+        inputs = [torch.randn(shape, device="cuda"), new_shape]
+
+        def fusion_func(fd: FusionDefinition):
+            t0 = fd.from_pytorch(inputs[0])
+            n_shape = fd.define_vector(10)
+
+            t1 = fd.ops.reshape(t0, n_shape)
+            t2 = fd.ops.sum(t1, axes=[3])
+
+            fd.add_output(t2)
+
+        nvf_out, _ = self.exec_nvfuser(fusion_func, inputs)
+        eager_out = torch.sum(inputs[0].reshape(new_shape), dim=3)
+        self.assertEqual(eager_out, nvf_out[0])
+
     # Testing a scenario where a broadcast requires a symbolic output shape
     def test_tensor_shape(self):
         inputs = [

test/test_gpu_fused_reduction.cpp

@@ -2559,4 +2559,29 @@ TEST_F(NVFuserTest, FusionCrossEntropyGatherPattern_CUDA) {
   testValidate(&fusion, cg_outputs, inputs, {ref}, __LINE__, __FILE__);
 }
 
+TEST_F(NVFuserTest, FusionTensorRankLimit) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  std::vector<int64_t> input_shape;
+  for (__attribute__((unused)) auto i : c10::irange(12)) {
+    input_shape.push_back(3);
+  }
+
+  auto tv0 = makeSymbolicTensor(input_shape.size());
+  fusion->addInput(tv0);
+  auto tv1 = sum(tv0, {3});
+  fusion->addOutput(tv1);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  at::Tensor t0 = at::randn(input_shape, options);
+  std::vector<c10::IValue> aten_inputs({t0});
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  auto cg_outputs = executor_cache.runFusionWithInputs(aten_inputs);
+
+  testValidate(
+      executor_cache.fusion(), cg_outputs, aten_inputs, __LINE__, __FILE__);
+}
+
 } // namespace nvfuser