Track variable-length tensor sizes through compilation

sdm/tensor/var_len.py

@@ -71,6 +71,7 @@ class VarLenTensor(Tensor):
     _data: Tensor
     _offset: Tensor
     _valid: Tensor | None
+    _layout: Tensor
 
     # Constructors ############################################################
 
@@ -248,6 +249,11 @@ def _make_var_len_tensor(
         out._data = data
         out._offset = offset
         out._valid = valid
+        # Carry size and stride as an empty inner tensor so torch.compile
+        # tracks them; outer sizes of nested subclasses can go stale.
+        out._layout = offset.new_empty_strided(
+            (*size, 0), (*stride, 1), dtype=torch.uint8
+        )
         return out
 
     @classmethod
@@ -555,16 +561,10 @@ def decorator(my_function: Callable[..., Any]) -> Callable[..., Any]:
 
     def __tensor_flatten__(self) -> tuple[list[str], tuple[Any, ...]]:
         # Offsets first avoid a torch.compile symbolic-shape check failure.
-        attrs = ["_offset", "_data"]
+        attrs = ["_offset", "_data", "_layout"]
         if self._valid is not None:
             attrs.append("_valid")
-        ctx = (
-            self.__class__,
-            self.storage_offset(),
-            tuple(self.size()),
-            tuple(self.stride()),
-        )
-        return attrs, ctx
+        return attrs, (self.__class__, self.storage_offset())
 
     @staticmethod
     def __tensor_unflatten__(
@@ -573,15 +573,18 @@ def __tensor_unflatten__(
         outer_size: tuple[int, ...],
         outer_stride: tuple[int, ...],
     ) -> VarLenTensor:
-        cls, storage_offset, _, _ = ctx
-        return cls(
+        cls, storage_offset = ctx
+        layout = inner_tensors["_layout"]
+        out = cls(
             data=inner_tensors["_data"],
             offset=inner_tensors["_offset"],
             valid=inner_tensors.get("_valid"),
-            size=outer_size,
-            stride=outer_stride,
+            size=layout.size()[:-1],
+            stride=layout.stride()[:-1],
             storage_offset=storage_offset,
         )
+        out._layout = layout
+        return out
 
     def __reduce_ex__(self, proto: SupportsIndex) -> Any:
         args = (

test/tensor/test_categorical.py

@@ -7,6 +7,7 @@
 import pyarrow as pa
 import pytest
 import torch
+from torch._dynamo.testing import CompileCounterWithBackend
 
 from sdm import CategoricalTensor, ColumnarTensor, NullableTensor, StringTensor
 from sdm.testing import onlyCUDA, withCUDA
@@ -676,3 +677,29 @@ def test_isnan() -> None:
     assert tensor.isnan().equal(expected)
     assert torch.isfinite(tensor).equal(~expected)
     assert tensor.isfinite().equal(~expected)
+
+
+@pytest.mark.parametrize("dynamic", [None, True])
+def test_compiled_string_categories_follow_input_sizes(
+    compilation: None,
+    dynamic: bool | None,
+) -> None:
+    def make(rows: int, size: int) -> CategoricalTensor:
+        categories = StringTensor.from_list([f"c{i}" for i in range(size)])
+        code = (torch.arange(rows) % size).view(rows, 1)
+        return CategoricalTensor(code, (categories,))
+
+    def fn(
+        x: CategoricalTensor, y: CategoricalTensor
+    ) -> tuple[torch.Tensor, ...]:
+        x, y = (cast(CategoricalTensor, t.unsqueeze(0)) for t in (x, y))
+        return x.categories[0], y.categories[0]
+
+    counter = CompileCounterWithBackend("aot_eager")
+    compiled = torch.compile(fn, backend=counter, dynamic=dynamic)
+    for rows, size in [(8, 4), (8, 4), (5, 3), (12, 6), (7, 5)]:
+        x, y = make(rows, size), make(rows + 1, size - 1)
+        for out, inp in zip(compiled(x, y), (x, y), strict=True):
+            assert out.tolist() == inp.categories[0].tolist()
+    # New row counts and vocabulary sizes reuse one dynamic graph:
+    assert counter.frame_count <= 2