[Tripy] Eliminate need for skip_num_stack_entries argument in convert_to_tensors (#333)

Addresses issue #310. The only use of `skip_num_stack_entries` was for
`slice_helper` and addressing this issue in a systematic manner would
likely require building in many hacks and assumptions, so the approach
here is just to manually override the stack information in that one
function.

Author: slyubomirsky

tripy/tripy/frontend/trace/ops/slice.py

@@ -250,8 +250,37 @@ def clamp_bound(bound: Union[int, Tensor]) -> Union[int, Tensor]:
     return out
 
 
-# Because the helper is called inside another function, we need to skip one entry in the call stack to find
-# the original call to user code.
-@frontend_utils.convert_to_tensors(skip_num_stack_entries=1)
+@frontend_utils.convert_to_tensors()
 def slice_helper(tensor, *slice_params: TensorLike):
+    from tripy.utils import get_arg_candidate_column_offsets
+
+    # The default behavior of convert_to_tensors will not add the correct column info to the slice params
+    # because this call occurs *inside* the overridden call to __getitem__, so we adjust the column info manually.
+
+    # Look for the stack frame index to __getitem__. We need to go one stack frame beyond to get to the *user* call of __getitem__.
+    # It will be the same for all the slice params
+    frame_index = -1
+    assert slice_params
+
+    for idx, source_info in enumerate(slice_params[0].stack_info):
+        if source_info._dispatch_target == "__getitem__":
+            frame_index = idx + 1
+            break
+
+    # convert_to_tensors should have taken care of this for us
+    assert frame_index >= 0, "No call to the __getitem__ dispatch found"
+
+    arg_names = ["tensor"] + ["slice_params"] * len(slice_params)
+    for arg_index, arg in enumerate(slice_params):
+        source_info = arg.stack_info[frame_index]
+
+        # Note: arg_index does not account for the positional arg, hence we add 1 for the index argument
+        candidates = get_arg_candidate_column_offsets(
+            source_info.code, 1 + arg_index, 1, "__getitem__", False, arg_names
+        )
+
+        # Now we can set the column range correctly
+        if len(candidates) == 1:
+            source_info.column_range = candidates[0]
+
     return Slice.build(inputs=[tensor, *slice_params])

tripy/tripy/frontend/utils.py

@@ -70,7 +70,8 @@ def empty_buffer():
 
 
 # Try to include correct column offsets for non-tensor arguments.
-def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, skip_num_stack_entries, arg_names):
+def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, arg_names):
+    from tripy import function_registry
     from tripy.frontend.tensor import Tensor
 
     assert isinstance(arg, Tensor), f"This function should only be called for objects that are already Tensor instances"
@@ -90,10 +91,10 @@ def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, skip_n
     # Find the first caller of this function that is NOT the function registry.
     # Also save the last dispatch target we see.
     dispatch_target = None
-    for idx, source_info in enumerate(arg.stack_info[WRAPPER_STACK_DEPTH + skip_num_stack_entries :]):
+    for idx, source_info in enumerate(arg.stack_info[WRAPPER_STACK_DEPTH:]):
         dispatch_target = source_info._dispatch_target or dispatch_target
         if source_info.module not in utils.get_module_names_to_exclude_from_stack_info():
-            frame_index = idx + WRAPPER_STACK_DEPTH + skip_num_stack_entries
+            frame_index = idx + WRAPPER_STACK_DEPTH
             break
     else:
         # Fallback path is just to look at the user code
@@ -118,12 +119,6 @@ def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, skip_n
         arg_index = 0 if arg_index == 1 else 1
         dispatch_target = dispatch_target.replace("__r", "__")
 
-    # Special case for __getitem__: It is variadic. Argument 0 is the tensor,
-    # and all subsequent arguments are slice parameters (in start, stop, step order).
-    # Hence, we subtract one to get the index of the slice parameters
-    if dispatch_target == "__getitem__":
-        arg_index -= 1
-
     candidates = utils.get_arg_candidate_column_offsets(
         source_info.code, arg_index, num_positional, dispatch_target or func_name, is_kwarg, arg_names
     )
@@ -136,9 +131,7 @@ def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, skip_n
 
 # NOTE: Conversion to tensors needs to be done via a decorator so that we can add stack information
 # for non-tensors. Without having full context of the function signature, it is otherwise difficult to do so.
-def convert_to_tensors(
-    targets: Set[str] = None, skip_num_stack_entries: int = 0, preprocess_args: Optional[Callable] = None
-):
+def convert_to_tensors(targets: Set[str] = None, preprocess_args: Optional[Callable] = None):
     """
     Decorator that converts specified arguments to Tensors or DimensionSizes.
     If the argument can be converted to a DimensionSize, it is. Otherwise, it is
@@ -152,17 +145,6 @@ def convert_to_tensors(
         targets: Names of arguments to convert to tensors. If not supplied, any arguments annotated
             with `TensorLike` or `ShapeLike` are converted.
 
-        skip_num_stack_entries: If the decorator is used on a function that is *called by*
-            a function that the user invokes, it will be necessary to skip stack entries
-            in order to get the column info from the user code. The number of entries skipped
-            should be equal to the nesting depth from a function called by user code
-            (if the decorated function is called by the user the depth is 0;
-            if the decorated function is called from a user function, the depth is 1; etc.).
-
-            NOTE: When using this, make sure any extra arguments to the decorated function are
-            passed as keyword arguments. Otherwise, the logic for determining column information
-            will break.
-
         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
@@ -242,7 +224,6 @@ def add_arg(arg):
                         name in kwargs,
                         len(args),
                         func.__name__,
-                        skip_num_stack_entries,
                         [name for name, _ in all_args],
                     )
 

tripy/tripy/utils/ast.py

@@ -139,12 +139,13 @@ def index_into_expr(node: ast.expr, index: int) -> ast.expr:
                 return node
 
             # If we have multiple dimensions specified, then we have a tuple of slices.
-            # Indices are given in as a list of start, stop, step
+            # NOTE: We subtract num_positional from the index because the slice arguments would
+            # be passed as *variadic arguments* to slice_helper and so would come after the positional argument
             if isinstance(node.slice, ast.Tuple):
-                element = node.slice.elts[index // 3]
-                arg_node = index_into_expr(element, index % 3)
+                element = node.slice.elts[(index - num_positional) // 3]
+                arg_node = index_into_expr(element, (index - num_positional) % 3)
             else:
-                arg_node = index_into_expr(node.slice, index)
+                arg_node = index_into_expr(node.slice, (index - num_positional))
 
         if arg_node is not None:
             candidates.append((indentation + arg_node.col_offset, indentation + arg_node.end_col_offset))