[Tripy] Handle variadic arguments in preprocess_args for the convert_to_tensors decorator (#329)

Addresses issue #311.

Author: slyubomirsky

tripy/tests/frontend/test_utils.py

@@ -178,3 +178,17 @@ def func(a: tp.types.TensorLike, b: tp.types.TensorLike):
         a, b = func(1, 2)
 
         assert b.tolist() == 3
+
+    def test_variadic_args(self):
+
+        def increment(a, *args):
+            return {"a": a + 1, "args": list(map(lambda arg: arg + 1, args))}
+
+        @convert_to_tensors(preprocess_args=increment)
+        def func(a: tp.Tensor, *args):
+            return [a] + list(args)
+
+        a, b, c = func(tp.Tensor(1), tp.Tensor(2), tp.Tensor(3))
+        assert a.tolist() == 2
+        assert b.tolist() == 3
+        assert c.tolist() == 4

tripy/tripy/backend/api/compile.py

@@ -151,7 +151,8 @@ def process_arg(name, arg):
         return arg
 
     new_args = []
-    for name, arg in utils.get_positional_arg_names(func, *args):
+    positional_arg_info, _ = utils.get_positional_arg_names(func, *args)
+    for name, arg in positional_arg_info:
         new_args.append(process_arg(name, arg))
 
     new_kwargs = {}

tripy/tripy/constraints.py

@@ -105,7 +105,7 @@ def wrapper(*args, **kwargs):
                 from tripy.common.datatype import dtype
                 from tripy.frontend.tensor import Tensor
 
-                merged_args = utils.merge_function_arguments(func, *args, **kwargs)
+                merged_args, _ = utils.merge_function_arguments(func, *args, **kwargs)
 
                 # The first arguments seen for each type variable. Other arguments with the same variable
                 # must use the same data types.

tripy/tripy/frontend/utils.py

@@ -18,7 +18,7 @@
 import functools
 import inspect
 from collections import deque
-from typing import Callable, List, Optional, Set, Union
+from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 from tripy import utils
 from tripy.common.exception import raise_error
@@ -166,7 +166,8 @@ def convert_to_tensors(
         preprocess_args: A callback used to preprocess arguments before potential conversion. If provided,
             this is always called, regardless of whether the decorator actually needed to perform conversion.
             This will be called with all arguments that were passed to the decorated function and should
-            return a dictionary of all updated arguments.
+            return a dictionary of all updated arguments. For a variadic arg, the dictionary entry for the name
+            should have a list of all the updated values.
     """
 
     def impl(func):
@@ -193,17 +194,20 @@ def wrapper(*args, **kwargs):
             from tripy.frontend.tensor import Tensor
             from tripy.frontend.trace.ops.cast import cast
 
-            all_args = utils.merge_function_arguments(func, *args, **kwargs)
+            all_args, var_arg_info = utils.merge_function_arguments(func, *args, **kwargs)
 
             if preprocess_args is not None:
+
+                var_arg_name, var_arg_start_idx = utils.default(var_arg_info, (None, None))
                 new_args = preprocess_args(*args, **kwargs)
-                # TODO (#311): Make this work for variadic arguments. If `name` appears multiple times in `all_args`, then
-                # we know we're dealing with a variadic argument. In that case, we could expect a list in `new_args` and
-                # then unpack it over the corresponding arguments in `all_args`.
                 for index in range(len(all_args)):
                     name, _ = all_args[index]
                     if name in new_args:
-                        all_args[index] = (name, new_args[name])
+                        if name == var_arg_name:
+                            assert var_arg_start_idx is not None
+                            all_args[index] = (name, new_args[name][index - var_arg_start_idx])
+                        else:
+                            all_args[index] = (name, new_args[name])
 
             # Materialize type variables from tensors.
             type_vars = {}

tripy/tripy/utils/utils.py

@@ -24,7 +24,7 @@
 import math
 import time
 import typing
-from typing import Any, List, Sequence, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union
 
 from colored import Fore, Style
 
@@ -415,10 +415,11 @@ def gen_uid(inputs=None, outputs=None):
 ##
 ## Functions
 ##
-def get_positional_arg_names(func, *args):
+def get_positional_arg_names(func, *args) -> Tuple[List[Tuple[str, Any]], Optional[Tuple[str, int]]]:
     # Returns the names of positional arguments by inspecting the function signature.
     # In the case of variadic positional arguments, we cannot determine names, so we use
-    # None instead.
+    # None instead. To assist in further processing, this function also returns the name
+    # and start index of the variadic args in a pair if present (None if not).
     signature = inspect.signature(func)
     arg_names = []
     varargs_name = None
@@ -432,12 +433,15 @@ def get_positional_arg_names(func, *args):
         arg_names.append(name)
 
     # For all variadic positional arguments, assign the name of the variadic group.
-    arg_names.extend([varargs_name] * (len(args) - len(arg_names)))
-    return list(zip(arg_names, args))
+    num_variadic_args = len(args) - len(arg_names)
+    variadic_start_idx = len(arg_names)
+    arg_names.extend([varargs_name] * num_variadic_args)
+    return list(zip(arg_names, args)), (varargs_name, variadic_start_idx) if num_variadic_args > 0 else None
 
 
-def merge_function_arguments(func, *args, **kwargs):
+def merge_function_arguments(func, *args, **kwargs) -> Tuple[List[Tuple[str, Any]], Optional[Tuple[str, int]]]:
     # Merge positional and keyword arguments, trying to determine names where possible.
-    all_args = get_positional_arg_names(func, *args)
+    # Also returns a pair containing the variadic arg name and start index if present (None otherwise).
+    all_args, var_arg_info = get_positional_arg_names(func, *args)
     all_args.extend(kwargs.items())
-    return all_args
+    return all_args, var_arg_info