diff --git a/sdm/_memory.py b/sdm/_memory.py index f9688ce10..ae55f6da4 100644 --- a/sdm/_memory.py +++ b/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( diff --git a/sdm/processing/numerical/rank_gaussian.py b/sdm/processing/numerical/rank_gaussian.py index fd95a19eb..8646808d4 100644 --- a/sdm/processing/numerical/rank_gaussian.py +++ b/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( diff --git a/sdm/processing/numerical/robust_scale.py b/sdm/processing/numerical/robust_scale.py index 5edc2ff83..a7b796604 100644 --- a/sdm/processing/numerical/robust_scale.py +++ b/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, )