From 2e329a2c91363d014d2e0e4d266f94d875aeec78 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:28:06 -0700
Subject: [PATCH 1/6] Track variable-length tensor sizes through compilation

Compiled code rebuilt nested string tensors (such as categorical
vocabularies) with the size seen at trace time, so a later call could
return a four-entry vocabulary with three entries. Guarding size and
stride by value avoided that but recompiled for every new row count or
vocabulary size. Carry size and stride in an empty inner tensor instead,
whose shape torch.compile tracks symbolically: new sizes now reuse one
dynamic graph and keep their exact contents.
---
 sdm/tensor/var_len.py           | 27 +++++++++++++++------------
 test/tensor/test_categorical.py | 27 +++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index a3c439b9c4..96f32571db 100644
--- a/sdm/tensor/var_len.py
+++ b/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 = (
diff --git a/test/tensor/test_categorical.py b/test/tensor/test_categorical.py
index 7709fec21c..e1780770b2 100644
--- a/test/tensor/test_categorical.py
+++ b/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

From e7801a676ca68ae4eb7449fb56831c57cb7f4d1b Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:37:45 -0700
Subject: [PATCH 2/6] Clarify why sizes are kept in an inner tensor

---
 sdm/tensor/var_len.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index 96f32571db..749b49e2f7 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -249,8 +249,8 @@ 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.
+        # torch.compile tracks inner tensor shapes but not the size of a
+        # nested subclass (e.g. a vocabulary), so keep size in a shape.
         out._layout = offset.new_empty_strided(
             (*size, 0), (*stride, 1), dtype=torch.uint8
         )

From 66e0049f93491bbff7ea8b796c92ea82848b4dfe Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:41:07 -0700
Subject: [PATCH 3/6] Say why sizes are mirrored in an empty tensor

---
 sdm/tensor/var_len.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index 749b49e2f7..c048e6d7ea 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -249,8 +249,8 @@ def _make_var_len_tensor(
         out._data = data
         out._offset = offset
         out._valid = valid
-        # torch.compile tracks inner tensor shapes but not the size of a
-        # nested subclass (e.g. a vocabulary), so keep size in a shape.
+        # Mirror size and stride in an empty tensor's shape so torch.compile
+        # can make them symbolic (it fixes a nested subclass's own size).
         out._layout = offset.new_empty_strided(
             (*size, 0), (*stride, 1), dtype=torch.uint8
         )

From 7d37043be0ac1990528bda5fc530e4cb8797c4b7 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:44:43 -0700
Subject: [PATCH 4/6] Explain what goes wrong without the mirrored size

---
 sdm/tensor/var_len.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index c048e6d7ea..a3203beca7 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -250,7 +250,8 @@ def _make_var_len_tensor(
         out._offset = offset
         out._valid = valid
         # Mirror size and stride in an empty tensor's shape so torch.compile
-        # can make them symbolic (it fixes a nested subclass's own size).
+        # can make them symbolic. A tensor nested in another one, like a
+        # categorical vocabulary, would otherwise keep its first traced size.
         out._layout = offset.new_empty_strided(
             (*size, 0), (*stride, 1), dtype=torch.uint8
         )

From ebedd8f9792d0937b4b6ee76534db5f26a09c0c0 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:46:08 -0700
Subject: [PATCH 5/6] Check compiled strided layouts and pass layout shape by
 keyword

---
 sdm/tensor/var_len.py       |  2 +-
 test/tensor/test_var_len.py | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index a3203beca7..fae8102ef2 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -253,7 +253,7 @@ def _make_var_len_tensor(
         # can make them symbolic. A tensor nested in another one, like a
         # categorical vocabulary, would otherwise keep its first traced size.
         out._layout = offset.new_empty_strided(
-            (*size, 0), (*stride, 1), dtype=torch.uint8
+            size=(*size, 0), stride=(*stride, 1), dtype=torch.uint8
         )
         return out
 
diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py
index 04c6f641b4..cd7c2bcafd 100644
--- a/test/tensor/test_var_len.py
+++ b/test/tensor/test_var_len.py
@@ -33,6 +33,16 @@ def reconstruct(inp: VarLenTensor) -> VarLenTensor:
     assert out.tolist() == tensor.tolist()
 
 
+def test_compiled_strided_layout(compilation: None) -> None:
+    rows = [["a", "bc", "d"], ["e", "", "fg"]]
+    tensor = StringTensor.from_list(rows).t()[1:]  # 2D, strided, offset
+    out = torch.compile(lambda x: x.detach(), backend="aot_eager")(tensor)
+    assert out.size() == tensor.size()
+    assert out.stride() == tensor.stride()
+    assert out.storage_offset() == tensor.storage_offset()
+    assert out.tolist() == tensor.tolist()
+
+
 def test_dtype_conversion() -> None:
     tensor = VarLenTensor.from_tensor(torch.arange(4).view(2, 2))
     assert repr(tensor) == "VarLenTensor(size=(2, 2), dtype=torch.int64)"

From bcd46f6038ea93ba05a26e59de8d2adaa8f2cd8f Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:50:46 -0700
Subject: [PATCH 6/6] Drop the separate strided layout test

---
 test/tensor/test_var_len.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py
index cd7c2bcafd..04c6f641b4 100644
--- a/test/tensor/test_var_len.py
+++ b/test/tensor/test_var_len.py
@@ -33,16 +33,6 @@ def reconstruct(inp: VarLenTensor) -> VarLenTensor:
     assert out.tolist() == tensor.tolist()
 
 
-def test_compiled_strided_layout(compilation: None) -> None:
-    rows = [["a", "bc", "d"], ["e", "", "fg"]]
-    tensor = StringTensor.from_list(rows).t()[1:]  # 2D, strided, offset
-    out = torch.compile(lambda x: x.detach(), backend="aot_eager")(tensor)
-    assert out.size() == tensor.size()
-    assert out.stride() == tensor.stride()
-    assert out.storage_offset() == tensor.storage_offset()
-    assert out.tolist() == tensor.tolist()
-
-
 def test_dtype_conversion() -> None:
     tensor = VarLenTensor.from_tensor(torch.arange(4).view(2, 2))
     assert repr(tensor) == "VarLenTensor(size=(2, 2), dtype=torch.int64)"