From 408955ff58b23e1ee4c57a60d731ab084e697bd7 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:29:10 -0700 Subject: [PATCH 1/3] Preserve bounded symbolic dimensions in preprocessing splits --- sdm/_memory.py | 10 ++++++ sdm/processing/numerical/rank_gaussian.py | 8 ++--- sdm/processing/numerical/robust_scale.py | 4 +-- test/test_memory.py | 38 +++++++++++++++++++++++ 4 files changed, 54 insertions(+), 6 deletions(-) create mode 100644 test/test_memory.py diff --git a/sdm/_memory.py b/sdm/_memory.py index f9688ce10f..ae55f6da48 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 fd95a19eb0..8646808d47 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 5edc2ff83a..a7b7966042 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, ) diff --git a/test/test_memory.py b/test/test_memory.py new file mode 100644 index 0000000000..4dd01e0d8a --- /dev/null +++ b/test/test_memory.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm._memory import _split_tensor + + +def _square_selected(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + selected = values[:, mask] + torch._check(selected.size(-1) <= values.size(-1)) + return torch.cat( + [chunk.square() for chunk in _split_tensor(selected, 8, dim=-1)], + dim=-1, + ) + + +def test_bounded_symbolic_split_compiles(compilation: None) -> None: + values = torch.arange(24.0).reshape(3, 8) + with torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True): + compiled = torch.compile( + _square_selected, backend="inductor", fullgraph=True + ) + for width in (0, 1, 3, 8): + mask = torch.arange(8) < width + torch.testing.assert_close( + compiled(values, mask), values[:, mask].square() + ) + + +def test_split_preserves_chunk_boundaries() -> None: + values = torch.arange(8.0).reshape(1, 8) + actual = torch.cat( + [chunk.cumsum(dim=-1) for chunk in _split_tensor(values, 3, -1)], + dim=-1, + ) + expected = torch.tensor([[0.0, 1.0, 3.0, 3.0, 7.0, 12.0, 6.0, 13.0]]) + torch.testing.assert_close(actual, expected) From 46128934cf8cacba81b1f06493c6ed6dd63dfea5 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:34:21 -0700 Subject: [PATCH 2/3] Simplify bounded split regression tests --- test/test_memory.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/test/test_memory.py b/test/test_memory.py index 4dd01e0d8a..4375154df2 100644 --- a/test/test_memory.py +++ b/test/test_memory.py @@ -8,7 +8,6 @@ def _square_selected(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: selected = values[:, mask] - torch._check(selected.size(-1) <= values.size(-1)) return torch.cat( [chunk.square() for chunk in _split_tensor(selected, 8, dim=-1)], dim=-1, @@ -18,9 +17,7 @@ def _square_selected(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: def test_bounded_symbolic_split_compiles(compilation: None) -> None: values = torch.arange(24.0).reshape(3, 8) with torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True): - compiled = torch.compile( - _square_selected, backend="inductor", fullgraph=True - ) + compiled = torch.compile(_square_selected, fullgraph=True) for width in (0, 1, 3, 8): mask = torch.arange(8) < width torch.testing.assert_close( @@ -29,10 +26,7 @@ def test_bounded_symbolic_split_compiles(compilation: None) -> None: def test_split_preserves_chunk_boundaries() -> None: - values = torch.arange(8.0).reshape(1, 8) - actual = torch.cat( - [chunk.cumsum(dim=-1) for chunk in _split_tensor(values, 3, -1)], - dim=-1, - ) - expected = torch.tensor([[0.0, 1.0, 3.0, 3.0, 7.0, 12.0, 6.0, 13.0]]) - torch.testing.assert_close(actual, expected) + values = torch.arange(8) + chunks = _split_tensor(values, 3, dim=0) + expected = [[0, 1, 2], [3, 4, 5], [6, 7]] + assert [chunk.tolist() for chunk in chunks] == expected From bccd1ee7c1c455bb7a370be17976d4cc9708cf3f Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:35:50 -0700 Subject: [PATCH 3/3] Remove added symbolic split tests --- test/test_memory.py | 32 -------------------------------- 1 file changed, 32 deletions(-) delete mode 100644 test/test_memory.py diff --git a/test/test_memory.py b/test/test_memory.py deleted file mode 100644 index 4375154df2..0000000000 --- a/test/test_memory.py +++ /dev/null @@ -1,32 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import torch - -from sdm._memory import _split_tensor - - -def _square_selected(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - selected = values[:, mask] - return torch.cat( - [chunk.square() for chunk in _split_tensor(selected, 8, dim=-1)], - dim=-1, - ) - - -def test_bounded_symbolic_split_compiles(compilation: None) -> None: - values = torch.arange(24.0).reshape(3, 8) - with torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True): - compiled = torch.compile(_square_selected, fullgraph=True) - for width in (0, 1, 3, 8): - mask = torch.arange(8) < width - torch.testing.assert_close( - compiled(values, mask), values[:, mask].square() - ) - - -def test_split_preserves_chunk_boundaries() -> None: - values = torch.arange(8) - chunks = _split_tensor(values, 3, dim=0) - expected = [[0, 1, 2], [3, 4, 5], [6, 7]] - assert [chunk.tolist() for chunk in chunks] == expected