diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py index 8557d3845..39902348c 100644 --- a/sdm/tensor/var_len.py +++ b/sdm/tensor/var_len.py @@ -1577,3 +1577,49 @@ def _materialize( ) out._max_length = inp._max_length return out + + +def _select_entries(inp: VarLenTensor, index: Tensor) -> VarLenTensor: + """Compact selected entries into bounded backing storage. + + The input and index are one-dimensional. Padding after the final offset + keeps the output capacity static without truncating selected values. + """ + index = index.where(index >= 0, index + inp.numel()) + positions = index * inp.stride(0) + inp._logical_storage_offset + start = inp._offset[positions] + end = inp._offset[positions + 1] + max_length = inp._max_length + capacity = index.numel() * max_length + if capacity > torch.iinfo(inp._offset.dtype).max: + return cast(VarLenTensor, inp[index]) + offset = torch.cat((start.new_zeros(1), (end - start).cumsum(0))) + offset = offset.to(inp._offset.dtype) + if index.numel() == 0 or inp._data.numel() == 0: + data = inp._data.new_empty(0) + elif inp._data.device.type == "cpu": + byte = torch.arange(max_length, device=inp.device, dtype=offset.dtype) + used = byte < (end - start)[:, None] + source = (start[:, None] + byte).clamp_max(inp._data.numel() - 1) + target = torch.where(used, offset[:-1, None] + byte, capacity) + values = inp._data[source].where(used, 0) + # Padding writes only to the extra byte, outside the logical payload. + data = inp._data.new_zeros(capacity + 1).scatter_( + 0, target.flatten().long(), values.flatten() + )[:-1] + else: + byte = torch.arange(capacity, device=inp.device, dtype=offset.dtype) + owner = torch.searchsorted(offset[1:], byte, right=True) + owner = owner.clamp_max(index.numel() - 1) + source = start[owner] + byte - offset[owner] + used = byte < offset[-1] + source = source.where(used, 0) + data = inp._data[source].where(used, 0) + out = inp.__class__( + data=data, + offset=offset, + valid=inp._valid[positions] if inp._valid is not None else None, + size=index.shape, + ) + out._max_length = max_length + return out diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py index 864a4e085..f1be5030e 100644 --- a/test/tensor/test_var_len.py +++ b/test/tensor/test_var_len.py @@ -9,7 +9,7 @@ from torch import Tensor from sdm import StringTensor, VarLenTensor -from sdm.tensor.var_len import _compact +from sdm.tensor.var_len import _compact, _select_entries from sdm.testing import onlyCUDA, withCUDA @@ -1034,3 +1034,88 @@ def test_reshape_inference_mode() -> None: assert not torch.is_inference(view) assert torch.is_inference(copied) + + +@withCUDA +@pytest.mark.parametrize("offset_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("indices", [[], [2, 0, 2, 1], [0, 1, 2], [-1, 0]]) +@pytest.mark.parametrize("layout", ["slice", "stride", "broadcast"]) +def test_native_string_selection( + device: torch.device, + offset_dtype: torch.dtype, + indices: list[int], + layout: str, +) -> None: + values = StringTensor.from_arrow( + pa.array(["unused", "é🙂", None, "", "long尾"]) + ) + assert values._valid is not None + values = StringTensor( + data=values._data.to(device), + offset=values._offset.to(device=device, dtype=offset_dtype), + valid=values._valid.to(device), + size=values.shape, + ) + if layout == "slice": + values = cast(StringTensor, values[1:4]) + elif layout == "stride": + values = cast(StringTensor, values[::2]) + else: + values = cast(StringTensor, values[1:2].expand(3)) + index = torch.tensor(indices, device=device, dtype=torch.int64) + expected = cast(StringTensor, values[index]) + actual = _select_entries(values, index) + assert actual.tolist() == expected.tolist() + assert actual.device == expected.device + assert actual.data_offset[1].dtype == expected.data_offset[1].dtype + assert actual.data_offset[1].diff().equal(expected.data_offset[1].diff()) + + +def test_native_string_selection_fullgraph(compilation: None) -> None: + values = StringTensor.from_arrow(pa.array(["é🙂", None, "", "long尾"])) + index = torch.tensor([3, 0, 3, 1, 2]) + compiled = torch.compile(_select_entries, fullgraph=True) + actual = compiled(values, index) + assert actual.tolist() == values[index].tolist() + + +def test_native_string_selection_meta() -> None: + values = StringTensor( + data=torch.empty(12, dtype=torch.uint8, device="meta"), + offset=torch.empty(4, dtype=torch.int32, device="meta"), + valid=torch.empty(3, dtype=torch.bool, device="meta"), + size=(3,), + ) + index = torch.empty(4, dtype=torch.int64, device="meta") + actual = _select_entries(values, index) + assert actual.shape == index.shape + assert actual.dtype == torch.uint8 + assert actual.device == index.device + assert actual._offset.shape == (index.numel() + 1,) + assert actual._offset.dtype == torch.int32 + assert actual._data.numel() == index.numel() * 12 + assert actual.valid is not None + assert actual.valid.shape == index.shape + + +@pytest.mark.parametrize("empty", [False, True]) +def test_native_string_selection_empty_payload(empty: bool) -> None: + values = StringTensor.from_arrow( + pa.array([], type=pa.string()) if empty else pa.array(["", None, ""]) + ) + index = torch.tensor([] if empty else [2, 0, 1, 2], dtype=torch.int64) + actual = _select_entries(values, index) + assert actual.tolist() == values[index].tolist() + + +def test_native_string_selection_null_payload() -> None: + values = StringTensor( + data=torch.tensor([97, 255, 98], dtype=torch.uint8), + offset=torch.tensor([0, 1, 2, 3], dtype=torch.int32), + valid=torch.tensor([True, False, True]), + size=(3,), + ) + index = torch.tensor([1, 2, 0, 1]) + actual = _select_entries(values, index) + assert actual.tolist() == values[index].tolist() + assert actual.data_offset[1].diff().tolist() == [1, 1, 1, 1]