replacing named_dims with shared_dims

Author: apbose

py/torch_tensorrt/_Input.py

@@ -51,7 +51,7 @@ class _ShapeMode(Enum):
     torch_tensor: torch.Tensor = None
     name: str = ""
     is_shape_tensor: bool = False
-    name_dims: Dict[int, str] = (
+    shared_dims: Dict[int, str] = (
         {}
     )  #: Optional {axis_index: name} for dynamic axes. The same name across inputs is exported as one shared ``torch.export.Dim`` (e.g. a batch axis shared by ``input_ids`` and ``attention_mask``).
 
@@ -165,19 +165,19 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                         "if you try to run inference with empty tensor inputs."
                     )
 
-                if "name_dims" in kwargs and kwargs["name_dims"]:
-                    self.name_dims = Input._parse_name_dims(
-                        kwargs["name_dims"], self.shape
+                if "shared_dims" in kwargs and kwargs["shared_dims"]:
+                    self.shared_dims = Input._parse_shared_dims(
+                        kwargs["shared_dims"], self.shape
                     )
 
         else:
             raise ValueError(
                 f"Unexpected number of positional arguments for class Input \n    Found {len(args)} arguments, expected either zero or a single positional arguments"
             )
 
-        if kwargs.get("name_dims") and self.shape_mode != Input._ShapeMode.DYNAMIC:
+        if kwargs.get("shared_dims") and self.shape_mode != Input._ShapeMode.DYNAMIC:
             raise ValueError(
-                "name_dims is only valid for dynamic inputs (min_shape/opt_shape/max_shape); "
+                "shared_dims is only valid for dynamic inputs (min_shape/opt_shape/max_shape); "
                 "it has no meaning for a statically shaped Input."
             )
 
@@ -276,29 +276,29 @@ def equivalent_spec(a: Input, b: Input) -> bool:
             return all(checks)
 
     @staticmethod
-    def _parse_name_dims(
-        name_dims: Any, shape: Dict[str, Tuple[int, ...]]
+    def _parse_shared_dims(
+        shared_dims: Any, shape: Dict[str, Tuple[int, ...]]
     ) -> Dict[int, str]:
-        """Validate and normalize the ``name_dims`` mapping ({axis: name}).
+        """Validate and normalize the ``shared_dims`` mapping ({axis: name}).
 
         Each named axis must be a valid index into the shape and must be
         genuinely dynamic (``min != max``); a static axis cannot vary, so naming
         it for cross-input sharing is a user error.
         """
-        if not isinstance(name_dims, dict):
+        if not isinstance(shared_dims, dict):
             raise TypeError(
-                f"name_dims must be a dict of {{axis_index: name}}, got {type(name_dims)}"
+                f"shared_dims must be a dict of {{axis_index: name}}, got {type(shared_dims)}"
             )
         rank = len(shape["min_shape"])
         parsed: Dict[int, str] = {}
-        for axis, dim_name in name_dims.items():
+        for axis, dim_name in shared_dims.items():
             if not isinstance(axis, int) or not (0 <= axis < rank):
                 raise ValueError(
-                    f"name_dims key {axis!r} is not a valid axis index for an input of rank {rank}"
+                    f"shared_dims key {axis!r} is not a valid axis index for an input of rank {rank}"
                 )
             if not isinstance(dim_name, str) or not dim_name:
                 raise ValueError(
-                    f"name_dims value for axis {axis} must be a non-empty string, got {dim_name!r}"
+                    f"shared_dims value for axis {axis} must be a non-empty string, got {dim_name!r}"
                 )
             if shape["min_shape"][axis] == shape["max_shape"][axis]:
                 raise ValueError(

py/torch_tensorrt/dynamo/_tracer.py

@@ -73,7 +73,7 @@ def trace(
     device = to_torch_device(kwargs.get("device", default_device()))
     torch_arg_inputs = get_torch_inputs(arg_inputs, device)
     torch_kwarg_inputs = get_torch_inputs(kwarg_inputs, device)
-    # Build dynamic shapes from the Input objects. Inputs carrying name_dims
+    # Build dynamic shapes from the Input objects. Inputs carrying shared_dims
     # share a Dim across inputs via the registry; the rest get an independent
     # per-input Dim.
     dim_registry = build_dim_registry(arg_inputs, kwarg_inputs)
@@ -108,7 +108,7 @@ def _collect_inputs(obj: Any) -> list[Input]:
 
 
 def build_dim_registry(arg_inputs: Any, kwarg_inputs: Any) -> dict[str, Any]:
-    """Build a ``{name: torch.export.Dim}`` registry from Input.name_dims.
+    """Build a ``{name: torch.export.Dim}`` registry from Input.shared_dims.
 
     The same name appearing on multiple inputs yields a single shared ``Dim``
     instance, so ``torch.export`` treats those axes as one symbol. Conflicting
@@ -117,13 +117,13 @@ def build_dim_registry(arg_inputs: Any, kwarg_inputs: Any) -> dict[str, Any]:
     registry: dict[str, Any] = {}
     bounds: dict[str, tuple[int, int]] = {}
     for inp in _collect_inputs(arg_inputs) + _collect_inputs(kwarg_inputs):
-        name_dims = getattr(inp, "name_dims", None)
-        if not name_dims or inp.shape_mode != Input._ShapeMode.DYNAMIC:
+        shared_dims = getattr(inp, "shared_dims", None)
+        if not shared_dims or inp.shape_mode != Input._ShapeMode.DYNAMIC:
             continue
         assert isinstance(inp.shape, dict)
         min_shape = inp.shape["min_shape"]
         max_shape = inp.shape["max_shape"]
-        for axis, dim_name in name_dims.items():
+        for axis, dim_name in shared_dims.items():
             lo, hi = int(min_shape[axis]), int(max_shape[axis])
             if dim_name in bounds:
                 if bounds[dim_name] != (lo, hi):
@@ -184,15 +184,15 @@ def get_dynamic_shapes(
             min_shape = input.shape["min_shape"]
             opt_shape = input.shape["opt_shape"]
             max_shape = input.shape["max_shape"]
-            name_dims = getattr(input, "name_dims", None) or {}
+            shared_dims = getattr(input, "shared_dims", None) or {}
             assert len(min_shape) == len(opt_shape) == len(max_shape)
             for dim in range(len(min_shape)):
                 if min_shape[dim] == opt_shape[dim] == max_shape[dim]:
                     continue
-                elif dim_registry is not None and dim in name_dims:
+                elif dim_registry is not None and dim in shared_dims:
                     # Named axis: reuse the shared Dim so axes with the same
                     # name across inputs become a single exported symbol.
-                    dynamic_dims[dim] = dim_registry[name_dims[dim]]
+                    dynamic_dims[dim] = dim_registry[shared_dims[dim]]
                 else:
                     dynamic_dims[dim] = Dim(
                         input.name + "_" + str(dim),

tests/py/dynamo/models/test_shared_dynamic_dim.py

@@ -1,6 +1,6 @@
 # type: ignore
 """
-Tests for sharing a dynamic dimension across inputs via ``Input(name_dims=...)``.
+Tests for sharing a dynamic dimension across inputs via ``Input(shared_dims=...)``.
 
 Background: when a model takes multiple inputs whose dynamic axes must be
 **equal at runtime** (e.g. HF encoders with ``input_ids`` / ``attention_mask``
@@ -9,7 +9,7 @@
 constraint check for any forward() that broadcasts across those axes (here:
 ``embed(input_ids) * mask.unsqueeze(-1)``), raising ``ConstraintViolationError``.
 
-``Input(name_dims={axis: name})`` lets the caller tag a dynamic axis with a
+``Input(shared_dims={axis: name})`` lets the caller tag a dynamic axis with a
 name; the same name across inputs is exported as a single shared ``Dim``. All
 the dynamic-shape intent lives on the ``Input`` objects -- no separate
 ``dynamic_shapes`` argument and no ``torch.export`` knowledge required at the
@@ -54,14 +54,14 @@ def _named_input(name: str, seq: int = 16, batch_min: int = 1, batch_max: int =
         max_shape=(batch_max, seq),
         dtype=torch.int64,
         name=name,
-        name_dims={0: "B"},
+        shared_dims={0: "B"},
     )
 
 
 @pytest.mark.unit
 @pytest.mark.critical
-def test_name_dims_shared_batch_kwarg_inputs():
-    """Shared batch axis declared via ``Input(name_dims={0: "B"})`` on both
+def test_shared_dims_shared_batch_kwarg_inputs():
+    """Shared batch axis declared via ``Input(shared_dims={0: "B"})`` on both
     kwarg inputs -- same name => one exported symbol; engine matches eager."""
     model = _SharedBatchEncoder().eval().cuda()
 
@@ -91,12 +91,12 @@ def test_name_dims_shared_batch_kwarg_inputs():
         cos_sim = cosine_similarity(ref, out)
         assertions.assertTrue(
             cos_sim > COSINE_THRESHOLD,
-            f"name_dims shared batch (kwargs) out-of-tolerance at bs={bs}: cos_sim={cos_sim}",
+            f"shared_dims shared batch (kwargs) out-of-tolerance at bs={bs}: cos_sim={cos_sim}",
         )
 
 
 @pytest.mark.unit
-def test_name_dims_shared_batch_positional_inputs():
+def test_shared_dims_shared_batch_positional_inputs():
     """Same feature with positional ``inputs=[...]`` instead of kwargs."""
     model = _SharedBatchEncoder().eval().cuda()
 
@@ -125,12 +125,12 @@ def test_name_dims_shared_batch_positional_inputs():
         cos_sim = cosine_similarity(ref, out)
         assertions.assertTrue(
             cos_sim > COSINE_THRESHOLD,
-            f"name_dims shared batch (positional) out-of-tolerance at bs={bs}: cos_sim={cos_sim}",
+            f"shared_dims shared batch (positional) out-of-tolerance at bs={bs}: cos_sim={cos_sim}",
         )
 
 
 @pytest.mark.unit
-def test_name_dims_shared_batch_mixed_args_and_kwargs():
+def test_shared_dims_shared_batch_mixed_args_and_kwargs():
     """input_ids passed positionally, attention_mask as a kwarg; both share "B"."""
     model = _SharedBatchEncoder().eval().cuda()
 
@@ -155,12 +155,12 @@ def test_name_dims_shared_batch_mixed_args_and_kwargs():
         cos_sim = cosine_similarity(ref, out)
         assertions.assertTrue(
             cos_sim > COSINE_THRESHOLD,
-            f"name_dims shared batch (mixed) out-of-tolerance at bs={bs}: cos_sim={cos_sim}",
+            f"shared_dims shared batch (mixed) out-of-tolerance at bs={bs}: cos_sim={cos_sim}",
         )
 
 
 @pytest.mark.unit
-def test_name_dims_conflicting_ranges_raises():
+def test_shared_dims_conflicting_ranges_raises():
     """Same name with different (min, max) across inputs is a user error."""
     from torch_tensorrt.dynamo._tracer import build_dim_registry
 
@@ -172,23 +172,23 @@ def test_name_dims_conflicting_ranges_raises():
             max_shape=(4, seq),
             dtype=torch.int64,
             name="input_ids",
-            name_dims={0: "B"},
+            shared_dims={0: "B"},
         ),
         "attention_mask": torchtrt.Input(
             min_shape=(1, seq),
             opt_shape=(8, seq),
             max_shape=(8, seq),
             dtype=torch.int64,
             name="attention_mask",
-            name_dims={0: "B"},
+            shared_dims={0: "B"},
         ),
     }
     with assertions.assertRaises(ValueError):
         build_dim_registry((), inputs)
 
 
 @pytest.mark.unit
-def test_name_dims_rejected_on_static_axis():
+def test_shared_dims_rejected_on_static_axis():
     """Naming a static axis (min == max) is rejected at Input construction."""
     with assertions.assertRaises(ValueError):
         torchtrt.Input(
@@ -197,12 +197,12 @@ def test_name_dims_rejected_on_static_axis():
             max_shape=(1, 16),
             dtype=torch.int64,
             name="x",
-            name_dims={0: "B"},
+            shared_dims={0: "B"},
         )
 
 
 @pytest.mark.unit
-def test_name_dims_rejected_on_out_of_range_axis():
+def test_shared_dims_rejected_on_out_of_range_axis():
     """An axis index outside the input's rank is rejected at construction."""
     with assertions.assertRaises(ValueError):
         torchtrt.Input(
@@ -211,13 +211,13 @@ def test_name_dims_rejected_on_out_of_range_axis():
             max_shape=(4, 16),
             dtype=torch.int64,
             name="x",
-            name_dims={5: "B"},  # rank is 2; axis 5 does not exist
+            shared_dims={5: "B"},  # rank is 2; axis 5 does not exist
         )
 
 
 @pytest.mark.unit
 def test_default_path_unchanged_for_static_inputs():
-    """Sanity check: a fully static input with no name_dims is unchanged."""
+    """Sanity check: a fully static input with no shared_dims is unchanged."""
 
     class StaticModel(nn.Module):
         def __init__(self):