From 80fd0c37643a24cd35783a2c94f4d6267233485b Mon Sep 17 00:00:00 2001
From: Manuel Candales <mcandales@meta.com>
Date: Fri, 6 Feb 2026 16:53:26 -0500
Subject: [PATCH] Pass all length/ix arguments to AttributeEngine

---
 facto/inputgen/argument/engine.py     | 28 ++++++++++++++-------------
 test/inputgen/test_meta_arg_engine.py | 21 ++++++++++++++++++++
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/facto/inputgen/argument/engine.py b/facto/inputgen/argument/engine.py
index 476493f..d58357a 100644
--- a/facto/inputgen/argument/engine.py
+++ b/facto/inputgen/argument/engine.py
@@ -48,7 +48,7 @@ def hierarchy(argtype) -> List[Attribute]:
             return [Attribute.VALUE]
 
     def gen_structure_with_depth_and_length(
-        self, depth: int, length: int, focus: Attribute
+        self, args: tuple, depth: int, length: int, focus: Attribute
     ):
         if length == 0:
             yield ()
@@ -57,7 +57,7 @@ def gen_structure_with_depth_and_length(
         attr = self._hierarchy[-(depth + 1)]
 
         if attr in self.gen_list_mode:
-            yield from self.gen_structure_with_depth(depth, focus, length)
+            yield from self.gen_structure_with_depth(args, depth, focus, length)
             return
 
         focus_ixs = (
@@ -67,9 +67,13 @@ def gen_structure_with_depth_and_length(
             values = [()]
             for ix in range(length):
                 if ix == focus_ix:
-                    elements = self.gen_structure_with_depth(depth, focus, length, ix)
+                    elements = self.gen_structure_with_depth(
+                        args, depth, focus, length, ix
+                    )
                 else:
-                    elements = self.gen_structure_with_depth(depth, None, length, ix)
+                    elements = self.gen_structure_with_depth(
+                        args, depth, None, length, ix
+                    )
                 new_values = []
                 for elem in elements:
                     new_values += [t + (elem,) for t in values]
@@ -78,6 +82,7 @@ def gen_structure_with_depth_and_length(
 
     def gen_structure_with_depth(
         self,
+        args: tuple,
         depth: int,
         focus: Attribute,
         length: Optional[int] = None,
@@ -86,14 +91,9 @@ def gen_structure_with_depth(
         attr = self._hierarchy[-(depth + 1)]
 
         if ix is not None:
-            args = (self.deps, length, ix)
+            args += (length, ix)
         elif length is not None:
-            args = (
-                self.deps,
-                length,
-            )
-        else:
-            args = (self.deps,)
+            args += (length,)
 
         values = AttributeEngine(attr, self.constraints, self.valid, self.argtype).gen(
             focus, *args
@@ -103,11 +103,13 @@ def gen_structure_with_depth(
             if depth == 0:
                 yield v
             else:
-                yield from self.gen_structure_with_depth_and_length(depth - 1, v, focus)
+                yield from self.gen_structure_with_depth_and_length(
+                    args, depth - 1, v, focus
+                )
 
     def gen(self, focus: Attribute):
         depth = len(self._hierarchy) - 1
-        yield from self.gen_structure_with_depth(depth, focus)
+        yield from self.gen_structure_with_depth((self.deps,), depth, focus)
 
 
 class MetaArg:
diff --git a/test/inputgen/test_meta_arg_engine.py b/test/inputgen/test_meta_arg_engine.py
index d801441..30d7a1c 100644
--- a/test/inputgen/test_meta_arg_engine.py
+++ b/test/inputgen/test_meta_arg_engine.py
@@ -80,6 +80,27 @@ def test_optional_int(self):
         self.assertEqual(len(ms), 1)
         self.assertEqual(str(ms[0]), "ArgType.IntOpt None")
 
+    def test_tensor_list(self):
+        constraints = [
+            cp.Length.Eq(lambda deps: 4),
+            cp.Rank.Eq(lambda deps, length, ix: 2 if ix < 2 else 1),
+            cp.Size.Eq(lambda deps, length, ix, r, d: 32 if d == 0 else None),
+            cp.Size.Eq(
+                lambda deps, length, ix, r, d: 8 if ix == 0 and d == 1 else None
+            ),
+            cp.Size.Eq(
+                lambda deps, length, ix, r, d: 16 if ix == 1 and d == 1 else None
+            ),
+        ]
+        deps = []
+        outarg = False
+
+        engine = MetaArgEngine(outarg, ArgType.TensorList, constraints, deps, True)
+        ms = list(engine.gen(None))
+        self.assertEqual(len(ms), 1)
+        m = ms[0]
+        self.assertEqual(m.structure, ((32, 8), (32, 16), (32,), (32,)))
+
 
 if __name__ == "__main__":
     unittest.main()