improve tests, review fixes

Author: akhilg-nv

tripy/nvtripy/backend/api/compile.py

@@ -156,6 +156,7 @@ def add(a, b):
     trace_input_map = {}
     input_names = set()
     input_infos = {}
+    trace_inputs = []  # flattened list of trace input tensors in argument order
 
     # Set up names for the weights in the module to make the trace easier to read.
     if isinstance(func, Module):
@@ -184,6 +185,7 @@ def process_arg_input_info(name, arg):
 
             trace_input_map[name] = tensor
             input_names.add(name)
+            trace_inputs.append(tensor.trace_tensor)
 
             return tensor
 
@@ -199,35 +201,44 @@ def process_arg_input_info(name, arg):
 
             trace_input_map[name] = tensor
             input_names.add(name)
+            trace_inputs.append(tensor.trace_tensor)
 
             return tensor
 
         return arg
 
-    def process_arg(name, arg):
+    def process_arg_and_flag(name, arg):
         # Handle individual InputInfo or DimensionInputInfo objects
         if isinstance(arg, (InputInfo, DimensionInputInfo)):
-            return process_arg_input_info(name, arg)
+            return process_arg_input_info(name, arg), True
 
         # Handle containers of InputInfo objects
         if isinstance(arg, dict):
-            if any(isinstance(v, (InputInfo, DimensionInputInfo)) for v in arg.values()):
-                input_names.add(name)
-                result = {}
-                for key, value in arg.items():
-                    nested_name = f"{name}.{key}"
-                    result[key] = process_arg(nested_name, value)
-                return result
+            result = {}
+            has_input = False
+            for key, value in arg.items():
+                nested_name = f"{name}.{key}"
+                processed_child, child_has_input = process_arg_and_flag(nested_name, value)
+                result[key] = processed_child
+                has_input = has_input or child_has_input
+            return result, has_input
         elif isinstance(arg, (list, tuple)):
-            if any(isinstance(v, (InputInfo, DimensionInputInfo)) for v in arg):
-                input_names.add(name)
-                result = []
-                for idx, value in enumerate(arg):
-                    nested_name = f"{name}[{idx}]"
-                    result.append(process_arg(nested_name, value))
-                return type(arg)(result)
+            result_list = []
+            has_input = False
+            for idx, value in enumerate(arg):
+                nested_name = f"{name}[{idx}]"
+                processed_child, child_has_input = process_arg_and_flag(nested_name, value)
+                result_list.append(processed_child)
+                has_input = has_input or child_has_input
+            return type(arg)(result_list), has_input  # preserve sequence type
 
-        return arg
+        return arg, False
+
+    def process_arg(name, arg):
+        processed, has_input = process_arg_and_flag(name, arg)
+        if has_input:
+            input_names.add(name)
+        return processed
 
     compiled_arg_names = []
 
@@ -283,25 +294,7 @@ def process_arg(name, arg):
                 [f"Return value {index} was not a tensor: {repr(trace_out)}"],
             )
 
-    # Order of trace inputs also needs to match that of the compiled_arg_names
-    # For containers, we need to collect all individual trace tensors
-    def collect_trace_tensors(name):
-        """Collect trace tensors for a name, flattening containers."""
-        if name in trace_input_map:
-            # Regular InputInfo or DimensionInputInfo
-            return [trace_input_map[name].trace_tensor]
-        else:
-            # Collect all nested trace tensors inside the container
-            nested_tensors = []
-            for nested_name in sorted(trace_input_map.keys()):
-                if nested_name.startswith(f"{name}.") or nested_name.startswith(f"{name}["):
-                    nested_tensors.append(trace_input_map[nested_name].trace_tensor)
-            return nested_tensors
-
-    # Flatten all trace tensors from containers and individual inputs
-    trace_inputs = []
-    for name in compiled_arg_names:
-        trace_inputs.extend(collect_trace_tensors(name))
+    # We collected flattened trace inputs during traversal
     trace = Trace(
         [tensor.trace_tensor for tensor in trace_outputs],
         trace_inputs,

tripy/nvtripy/backend/api/executable.py

@@ -195,40 +195,44 @@ def add(a, b):
                 ],
             )
 
-        # Recursively extract inputs from containers to get individual tensors for validation and execution
+        # Recursively build a name->tensor map
         def extract_inputs(tensors, input_info_names):
-            def extract_recursive(tensor, name_prefix):
-                if isinstance(tensor, dict):
-                    result = []
-                    for key in sorted(tensor.keys()):
+            name_to_tensor = {}
+
+            def extract_recursive(value, name_prefix):
+                if name_prefix in input_info_names:
+                    name_to_tensor[name_prefix] = value
+                    return
+                if isinstance(value, dict):
+                    for key, item in value.items():
                         nested_name = f"{name_prefix}.{key}"
-                        if nested_name in input_info_names:
-                            result.append(tensor[key])
-                        else:
-                            result.extend(extract_recursive(tensor[key], nested_name))
-                    return result
-                elif isinstance(tensor, (list, tuple)):
-                    result = []
-                    for idx, value in enumerate(tensor):
+                        extract_recursive(item, nested_name)
+                elif isinstance(value, (list, tuple)):
+                    for idx, item in enumerate(value):
                         nested_name = f"{name_prefix}[{idx}]"
-                        if nested_name in input_info_names:
-                            result.append(value)
-                        else:
-                            result.extend(extract_recursive(value, nested_name))
-                    return result
-                else:  # Regular tensor
-                    if name_prefix in input_info_names:
-                        return [tensor]
-                    else:
-                        return []
-
-            flattened = []
+                        extract_recursive(item, nested_name)
+                else:
+                    print(f"Leaf tensor: {name_prefix}: {value}")
+                    return
+
             for name_idx, tensor in enumerate(tensors):
                 arg_name = self._arg_names[name_idx]
-                flattened.extend(extract_recursive(tensor, arg_name))
-            return flattened
+                extract_recursive(tensor, arg_name)
+
+            return name_to_tensor
 
-        flattened_tensors = extract_inputs(input_tensors, set(self.input_infos.keys()))
+        input_info_names = list(self.input_infos.keys())
+        name_to_tensor = extract_inputs(input_tensors, set(input_info_names))
+        try:
+            flattened_tensors = [name_to_tensor[name] for name in input_info_names]
+        except KeyError as missing:
+            raise_error(
+                f"Missing runtime tensor for input `{missing.args[0]}`.",
+                [
+                    "Ensure your provided containers include tensors for all compiled inputs.",
+                    f"Expected inputs: {input_info_names}",
+                ],
+            )
         expected_devices = ["gpu" if isinstance(info, InputInfo) else "cpu" for info in self.input_infos.values()]
 
         # Validate flattened tensors against input_infos

tripy/tests/backend/api/test_compile.py

@@ -242,44 +242,55 @@ def test_dimension_input(self):
             expected = (inp_cp + inp_cp).reshape((-1, reshape_dim))
             assert cp.array_equal(cp.from_dlpack(out), expected)
 
-    def test_compile_dict_input_info(self):
-        """Test compilation with dictionary of InputInfo objects."""
-
+    def test_compile_nested_dict_input_info(self):
         def func(data_dict):
-            return data_dict["a"] + data_dict["b"]
+            return data_dict["a"]["inner"] + data_dict["b"]["list"][0] + data_dict["b"]["list"][1]
 
         dict_input = {
-            "a": tp.InputInfo(shape=(2, 3), dtype=tp.float32),
-            "b": tp.InputInfo(shape=(2, 3), dtype=tp.float32),
+            "a": {
+                "inner": tp.InputInfo(shape=(2, 3), dtype=tp.float32),
+            },
+            "b": {
+                "list": [
+                    tp.InputInfo(shape=(2, 3), dtype=tp.float32),
+                    tp.InputInfo(shape=(2, 3), dtype=tp.float32),
+                ],
+            },
         }
         compiled_func = tp.compile(func, args=[dict_input])
 
-        test_dict = {"a": tp.ones((2, 3), dtype=tp.float32).eval(), "b": (tp.ones((2, 3), dtype=tp.float32) * 2).eval()}
+        test_dict = {
+            "a": {"inner": tp.ones((2, 3), dtype=tp.float32).eval()},
+            "b": {
+                "list": [
+                    (tp.ones((2, 3), dtype=tp.float32) * 2).eval(),
+                    (tp.ones((2, 3), dtype=tp.float32) * 3).eval(),
+                ]
+            },
+        }
         result = compiled_func(test_dict)
-        expected = test_dict["a"] + test_dict["b"]
+        expected = test_dict["a"]["inner"] + test_dict["b"]["list"][0] + test_dict["b"]["list"][1]
         assert cp.array_equal(cp.from_dlpack(result), cp.from_dlpack(expected))
 
-    def test_compile_nested_list_input_info(self):
-        """Test compilation with nested list containers."""
-
+    def test_compile_nested_sequence_input_info(self):
         def func(data_list):
             return data_list[0] + data_list[1][0] + data_list[1][1]
 
         list_input = [
             tp.InputInfo(shape=(2, 3), dtype=tp.float32),
-            [  # Nested list
+            [
                 tp.InputInfo(shape=(2, 3), dtype=tp.float32),
-                tp.ones((2, 3), dtype=tp.float32) * 2,  # Constant in nested list
+                tp.ones((2, 3), dtype=tp.float32) * 2,
             ],
         ]
         compiled_func = tp.compile(func, args=[list_input])
 
         test_list = [
             tp.ones((2, 3), dtype=tp.float32).eval(),
-            [  # Nested list in test data
+            (
                 (tp.ones((2, 3), dtype=tp.float32) * 3).eval(),
-                tp.ones((2, 3), dtype=tp.float32) * 2,  # Should match baked constant
-            ],
+                tp.ones((2, 3), dtype=tp.float32) * 2,
+            ),
         ]
         result = compiled_func(test_list)
         expected = test_list[0] + test_list[1][0] + test_list[1][1]
@@ -288,24 +299,35 @@ def func(data_list):
     def test_compile_mixed_containers_and_constants(self):
         """Test compilation with comprehensive mix: regular InputInfo, dict container, list container, and standalone constant."""
 
-        def func(regular_input, data_dict, data_list, constant_value):
-            return regular_input + data_dict["x"] + data_dict["y"] + data_list[0] + data_list[1] + constant_value
+        def func(regular_input, data_dict, data_list, const_in_dict, const):
+            return (
+                regular_input
+                + data_dict["x"]
+                + data_dict["y"]
+                + data_list[0]
+                + data_list[1]
+                + const_in_dict["z"]
+                + const
+            )
 
         regular_input = tp.InputInfo(shape=(2, 3), dtype=tp.float32)
         dict_input = {
             "x": tp.InputInfo(shape=(2, 3), dtype=tp.float32),
-            "y": tp.zeros((2, 3), dtype=tp.float32),  # Constant in dict
+            "y": tp.zeros((2, 3), dtype=tp.float32),
         }
         list_input = [tp.InputInfo(shape=(2, 3), dtype=tp.float32), tp.ones((2, 3), dtype=tp.float32) * 3]
-        constant_value = tp.ones((2, 3), dtype=tp.float32) * 5
+        const_in_dict = {"z": tp.ones((2, 3), dtype=tp.float32) * 5}
+        const = tp.ones((2, 3), dtype=tp.float32) * 6
 
-        compiled_func = tp.compile(func, args=[regular_input, dict_input, list_input, constant_value])
+        compiled_func = tp.compile(func, args=[regular_input, dict_input, list_input, const_in_dict, const])
 
         # Only InputInfo arguments should be in function signature
         test_regular = tp.ones((2, 3), dtype=tp.float32).eval()
         test_dict = {"x": (tp.ones((2, 3), dtype=tp.float32) * 2).eval(), "y": tp.zeros((2, 3), dtype=tp.float32)}
         test_list = [(tp.ones((2, 3), dtype=tp.float32) * 4).eval(), tp.ones((2, 3), dtype=tp.float32) * 3]
 
         result = compiled_func(test_regular, test_dict, test_list)
-        expected = test_regular + test_dict["x"] + test_dict["y"] + test_list[0] + test_list[1] + constant_value
+        expected = (
+            test_regular + test_dict["x"] + test_dict["y"] + test_list[0] + test_list[1] + const_in_dict["z"] + const
+        )
         assert cp.array_equal(cp.from_dlpack(result), cp.from_dlpack(expected))