Keep string length bounds out of compiled metadata

Author: aw471

sdm/tensor/string.py

@@ -35,8 +35,9 @@ class StringTensor(VarLenTensor):
         stride: The stride of the tensor.
         storage_offset: The offset into the logical ``offset`` storage.
         max_length: Upper bound on the byte length of every stored string,
-            or :obj:`None` if unknown. Lets compiled operations size buffers
+            or :obj:`None` if unknown. Lets eager operations size buffers
             without reading lengths from the device. Must not underestimate.
+            Compiled tensor operations see :obj:`None` and never recompile.
     """
 
     # NOTE Assume that `data` stores valid UTF-8 bytes and do not validate it.

sdm/tensor/var_len.py

@@ -63,8 +63,9 @@ class VarLenTensor(Tensor):
         stride: The stride of the tensor.
         storage_offset: The offset into the logical ``offset`` storage.
         max_length: Upper bound on the length of every stored entry, or
-            :obj:`None` if unknown. Lets compiled operations size buffers
+            :obj:`None` if unknown. Lets eager operations size buffers
             without reading lengths from the device. Must not underestimate.
+            Compiled tensor operations see :obj:`None` and never recompile.
     """
 
     ALLOWED_DTYPES: ClassVar[tuple[torch.dtype, ...] | None] = None
@@ -576,7 +577,9 @@ def __tensor_flatten__(self) -> tuple[list[str], tuple[Any, ...]]:
         attrs = ["_offset", "_data", "_layout"]
         if self._valid is not None:
             attrs.append("_valid")
-        return attrs, (self.__class__, self.storage_offset(), self._max_length)
+        # The length bound stays eager: compiled operations see ``None``.
+        # As context, each new query bound would compile another graph.
+        return attrs, (self.__class__, self.storage_offset())
 
     @staticmethod
     def __tensor_unflatten__(
@@ -585,7 +588,7 @@ def __tensor_unflatten__(
         outer_size: tuple[int, ...],
         outer_stride: tuple[int, ...],
     ) -> VarLenTensor:
-        cls, storage_offset, max_length = ctx
+        cls, storage_offset = ctx
         layout = inner_tensors["_layout"]
         out = cls(
             data=inner_tensors["_data"],
@@ -594,7 +597,6 @@ def __tensor_unflatten__(
             size=layout.size()[:-1],
             stride=layout.stride()[:-1],
             storage_offset=storage_offset,
-            max_length=max_length,
         )
         out._layout = layout
         return out

test/tensor/test_string.py

@@ -450,12 +450,15 @@ def test_bound_provenance(
         (values.clone(), texts),
         (values.detach(), texts),
         (values.to(device, copy=True), texts),
-        (compiled(values), texts),
         (pickle.loads(pickle.dumps(values)), texts),
     ):
         copied = cast(StringTensor, copied)
         assert copied._max_length == bound
         assert copied.tolist() == expected
+    # Compiled outputs carry no bound, so new bounds do not recompile.
+    compiled_values = cast(StringTensor, compiled(values))
+    assert compiled_values._max_length is None
+    assert compiled_values.tolist() == texts
     for suffix in ("", "x" * 16):
         extra = StringTensor.from_list([suffix], device=device)
         joined = cast(StringTensor, torch.cat((values, extra)))