Allow bounded single-chunk splits under torch.compile

sdm/_memory.py

@@ -4,6 +4,16 @@
 import os
 
 import torch
+from torch.fx.experimental.symbolic_shapes import statically_known_true
+
+
+def _split_tensor(
+    tensor: torch.Tensor, size: int, dim: int
+) -> tuple[torch.Tensor, ...]:
+    """Return one chunk when the dimension provably fits the split size."""
+    if statically_known_true(tensor.size(dim) <= size):
+        return (tensor,)
+    return tensor.split(size, dim=dim)
 
 
 def set_pinned_max_round_threshold(

sdm/processing/numerical/rank_gaussian.py

@@ -5,7 +5,7 @@
 from torch import Tensor
 
 from sdm import Stype, TableTensor
-from sdm._memory import split_size
+from sdm._memory import _split_tensor, split_size
 from sdm.processing import Processor
 from sdm.processing.numerical._stats import _isfinite, _ndtri
 from sdm.processing.numerical.quantile import _batched_interp
@@ -62,7 +62,7 @@ def _fit(
             device=numerical.device,
         )
         knots = []
-        for chunk in numerical.split(size, dim=-1):
+        for chunk in _split_tensor(numerical, size, dim=-1):
             num_rows = chunk.size(-2)
             max_knots = num_rows if self.max_knots is None else self.max_knots
             columns = chunk.to(dtype).movedim(-1, -2)  # [..., C, N]
@@ -152,8 +152,8 @@ def _transform(self, table: TableTensor) -> TableTensor:
             device=numerical.device,
         )
         for rows, out in zip(
-            columns.split(size, dim=-1),
-            output.split(size, dim=-2),
+            _split_tensor(columns, size, dim=-1),
+            _split_tensor(output, size, dim=-2),
             strict=True,
         ):
             quantiles = _batched_interp(

sdm/processing/numerical/robust_scale.py

@@ -5,7 +5,7 @@
 from torch import Tensor
 
 from sdm import Stype, TableTensor
-from sdm._memory import split_size
+from sdm._memory import _split_tensor, split_size
 from sdm.processing import InvertibleMixin, Processor
 from sdm.processing.numerical._stats import _isfinite
 
@@ -69,7 +69,7 @@ def _fit(
         lower, median, upper = torch.cat(
             [
                 chunk.nanquantile(q, dim=-2, keepdim=True)
-                for chunk in quantile_input.split(size, dim=-1)
+                for chunk in _split_tensor(quantile_input, size, dim=-1)
             ],
             dim=-1,
         )