From 7e724ce2e7ee8368a4e5613534465d8146e3c267 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Fri, 23 Oct 2026 09:35:47 +0000 Subject: [PATCH 1/2] Add categorical value decoding --- sdm/tensor/categorical.py | 40 +++++++++- test/tensor/test_categorical.py | 137 +++++++++++++++++++++++++++++++- 2 files changed, 175 insertions(+), 2 deletions(-) diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py index 2d21727627..c09d7d1ec2 100644 --- a/sdm/tensor/categorical.py +++ b/sdm/tensor/categorical.py @@ -14,7 +14,7 @@ from torch.utils import _pytree as pytree from typing_extensions import override -from sdm.tensor import StringTensor, VarLenTensor +from sdm.tensor import NullableTensor, StringTensor, VarLenTensor from sdm.tensor.io import ( arrow_as_tensor, combine_arrow_chunks, @@ -350,6 +350,44 @@ def categories(self) -> tuple[Tensor, ...]: r"""Return category vector for each categorical column.""" return self._categories + # Methods ################################################################# + + def decode(self) -> Tensor: + r"""Decode this tensor to its category values.""" + # Avoid a circular import through sdm.tensor. + from sdm.tensor import ColumnarTensor # noqa: PLC0415 + + if self.size(-1) == 0: + raise ValueError("Decoding requires at least one column") + + tensors = [] + for tensor in self.split(1, dim=-1): + code = tensor.code.squeeze(-1) + valid = code >= 0 + if not valid.all(): + tensor = tensor.categories[0][code.clamp(min=0)] + if isinstance(tensor, VarLenTensor): + data, offset = tensor.contiguous().data_offset + tensor = tensor.__class__( + data=data, + offset=offset, + valid=valid.reshape(-1), + size=tensor.size(), + ) + elif tensor.is_floating_point(): + tensor[~valid] = torch.nan + else: + tensor = NullableTensor(tensor, valid) + else: + tensor = tensor.categories[0][code] + + tensors.append(tensor) + + if len(tensors) == 1: + return tensors[0].unsqueeze(-1) + + return ColumnarTensor(tensors) + # Decorators ############################################################## @classmethod diff --git a/test/tensor/test_categorical.py b/test/tensor/test_categorical.py index 8596301dab..6a3885a7cc 100644 --- a/test/tensor/test_categorical.py +++ b/test/tensor/test_categorical.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import math import warnings from typing import cast @@ -8,7 +9,7 @@ import pytest import torch -from sdm import CategoricalTensor, StringTensor +from sdm import CategoricalTensor, ColumnarTensor, NullableTensor, StringTensor from sdm.testing import onlyCUDA, withCUDA @@ -74,6 +75,140 @@ def test_to_copy() -> None: assert out.dtype == torch.float32 +@withCUDA +@pytest.mark.parametrize("dtype", [torch.bool, torch.int64, torch.float32]) +@pytest.mark.parametrize("code_dtype", [torch.int32, torch.int64]) +def test_decode( + device: torch.device, dtype: torch.dtype, code_dtype: torch.dtype +) -> None: + categories = torch.tensor([1, 0], dtype=dtype, device=device) + code = torch.tensor( + [[[0], [1], [0]], [[1], [0], [1]]], dtype=code_dtype, device=device + ).transpose(0, 1) + tensor = CategoricalTensor(code=code, categories=(categories,)) + + out = tensor.decode() + + assert type(out) is torch.Tensor + assert out.device == device + assert out.dtype == dtype + expected = torch.tensor( + [[[1], [0]], [[0], [1]], [[1], [0]]], dtype=dtype, device=device + ) + torch.testing.assert_close(out, expected) + + +@withCUDA +def test_decode_strings(device: torch.device) -> None: + tensor = CategoricalTensor( + code=torch.tensor([[1], [0], [1]], device=device), + categories=(StringTensor.from_list(["a", "b"], device=device),), + ) + + out = tensor.decode() + + assert isinstance(out, StringTensor) + assert out.device == device + assert out.tolist() == [["b"], ["a"], ["b"]] + + +@withCUDA +def test_decode_empty(device: torch.device) -> None: + tensor = CategoricalTensor( + code=torch.empty(0, 1, dtype=torch.int64, device=device), + categories=(torch.empty(0, dtype=torch.bool, device=device),), + ) + + out = tensor.decode() + + assert out.shape == (0, 1) + assert out.dtype == torch.bool + assert out.device == device + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.bool, torch.int64, torch.float32]) +def test_decode_missing(device: torch.device, dtype: torch.dtype) -> None: + code = torch.tensor([[0, 1, -1], [1, -2, 0]], device=device).T.unsqueeze( + -1 + ) + tensor = CategoricalTensor( + code=code, + categories=(torch.tensor([1, 0], dtype=dtype, device=device),), + ) + + out = tensor.decode() + + assert out.shape == (3, 2, 1) + assert out.device == device + assert out.dtype == dtype + if dtype == torch.float32: + expected = torch.tensor( + [[[1], [0]], [[0], [torch.nan]], [[torch.nan], [1]]], + device=device, + ) + torch.testing.assert_close(out, expected, equal_nan=True) + else: + assert isinstance(out, NullableTensor) + assert out.tolist() == [[[1], [0]], [[0], [None]], [[None], [1]]] + + +@withCUDA +def test_decode_missing_strings(device: torch.device) -> None: + code = torch.tensor([[0, 1, -1], [1, -2, 0]], device=device).T.unsqueeze( + -1 + ) + tensor = CategoricalTensor( + code=code, + categories=(StringTensor.from_list(["a", "b"], device=device),), + ) + + out = tensor.decode() + + assert isinstance(out, StringTensor) + assert out.shape == (3, 2, 1) + assert out.device == device + assert out.tolist() == [[["a"], ["b"]], [["b"], [None]], [[None], ["a"]]] + + +@withCUDA +@pytest.mark.parametrize("missing", [False, True]) +def test_decode_columns(device: torch.device, missing: bool) -> None: + code = torch.tensor([[0, 1, 1], [1, 0, 0]], device=device) + if missing: + code[1] = -1 + tensor = CategoricalTensor( + code=code, + categories=( + torch.tensor([1.5, 2.5], device=device), + torch.tensor([False, True], device=device), + StringTensor.from_list(["a", "b"], device=device), + ), + ) + + out = tensor.decode() + + assert isinstance(out, ColumnarTensor) + assert out.shape == (2, 3) + assert out.device == device + if missing: + rows = out.tolist() + assert rows[0] == [1.5, True, "b"] + assert math.isnan(rows[1][0]) + assert rows[1][1:] == [None, None] + else: + assert out.tolist() == [[1.5, True, "b"], [2.5, False, "a"]] + + +def test_decode_requires_columns() -> None: + tensor = CategoricalTensor( + code=torch.empty(3, 0, dtype=torch.int64), + categories=(), + ) + with pytest.raises(ValueError, match="at least one column"): + tensor.decode() + + def test_from_arrow_string_values() -> None: tensor = CategoricalTensor.from_arrow( pa.array(["b", "a", None, "b"]), From 0e539c1c1d0280189aa52699d43f635611240b02 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Fri, 23 Oct 2026 09:40:19 +0000 Subject: [PATCH 2/2] update --- sdm/tensor/categorical.py | 1 - test/tensor/test_categorical.py | 135 ++++---------------------------- 2 files changed, 14 insertions(+), 122 deletions(-) diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py index c09d7d1ec2..2f7e56d249 100644 --- a/sdm/tensor/categorical.py +++ b/sdm/tensor/categorical.py @@ -354,7 +354,6 @@ def categories(self) -> tuple[Tensor, ...]: def decode(self) -> Tensor: r"""Decode this tensor to its category values.""" - # Avoid a circular import through sdm.tensor. from sdm.tensor import ColumnarTensor # noqa: PLC0415 if self.size(-1) == 0: diff --git a/test/tensor/test_categorical.py b/test/tensor/test_categorical.py index 6a3885a7cc..7709fec21c 100644 --- a/test/tensor/test_categorical.py +++ b/test/tensor/test_categorical.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import math import warnings from typing import cast @@ -76,137 +75,31 @@ def test_to_copy() -> None: @withCUDA -@pytest.mark.parametrize("dtype", [torch.bool, torch.int64, torch.float32]) -@pytest.mark.parametrize("code_dtype", [torch.int32, torch.int64]) -def test_decode( - device: torch.device, dtype: torch.dtype, code_dtype: torch.dtype -) -> None: - categories = torch.tensor([1, 0], dtype=dtype, device=device) - code = torch.tensor( - [[[0], [1], [0]], [[1], [0], [1]]], dtype=code_dtype, device=device - ).transpose(0, 1) - tensor = CategoricalTensor(code=code, categories=(categories,)) - - out = tensor.decode() - - assert type(out) is torch.Tensor - assert out.device == device - assert out.dtype == dtype - expected = torch.tensor( - [[[1], [0]], [[0], [1]], [[1], [0]]], dtype=dtype, device=device - ) - torch.testing.assert_close(out, expected) - - -@withCUDA -def test_decode_strings(device: torch.device) -> None: - tensor = CategoricalTensor( - code=torch.tensor([[1], [0], [1]], device=device), - categories=(StringTensor.from_list(["a", "b"], device=device),), - ) - - out = tensor.decode() - - assert isinstance(out, StringTensor) - assert out.device == device - assert out.tolist() == [["b"], ["a"], ["b"]] - - -@withCUDA -def test_decode_empty(device: torch.device) -> None: - tensor = CategoricalTensor( - code=torch.empty(0, 1, dtype=torch.int64, device=device), - categories=(torch.empty(0, dtype=torch.bool, device=device),), - ) - - out = tensor.decode() - - assert out.shape == (0, 1) - assert out.dtype == torch.bool - assert out.device == device - - -@withCUDA -@pytest.mark.parametrize("dtype", [torch.bool, torch.int64, torch.float32]) -def test_decode_missing(device: torch.device, dtype: torch.dtype) -> None: - code = torch.tensor([[0, 1, -1], [1, -2, 0]], device=device).T.unsqueeze( - -1 - ) - tensor = CategoricalTensor( - code=code, - categories=(torch.tensor([1, 0], dtype=dtype, device=device),), - ) - - out = tensor.decode() - - assert out.shape == (3, 2, 1) - assert out.device == device - assert out.dtype == dtype - if dtype == torch.float32: - expected = torch.tensor( - [[[1], [0]], [[0], [torch.nan]], [[torch.nan], [1]]], - device=device, - ) - torch.testing.assert_close(out, expected, equal_nan=True) - else: - assert isinstance(out, NullableTensor) - assert out.tolist() == [[[1], [0]], [[0], [None]], [[None], [1]]] - - -@withCUDA -def test_decode_missing_strings(device: torch.device) -> None: - code = torch.tensor([[0, 1, -1], [1, -2, 0]], device=device).T.unsqueeze( - -1 - ) +def test_decode(device: torch.device) -> None: tensor = CategoricalTensor( - code=code, - categories=(StringTensor.from_list(["a", "b"], device=device),), - ) - - out = tensor.decode() - - assert isinstance(out, StringTensor) - assert out.shape == (3, 2, 1) - assert out.device == device - assert out.tolist() == [[["a"], ["b"]], [["b"], [None]], [[None], ["a"]]] - - -@withCUDA -@pytest.mark.parametrize("missing", [False, True]) -def test_decode_columns(device: torch.device, missing: bool) -> None: - code = torch.tensor([[0, 1, 1], [1, 0, 0]], device=device) - if missing: - code[1] = -1 - tensor = CategoricalTensor( - code=code, + code=torch.tensor([[0, 1, 0], [-1, -1, -1], [1, 0, 1]], device=device), categories=( - torch.tensor([1.5, 2.5], device=device), - torch.tensor([False, True], device=device), StringTensor.from_list(["a", "b"], device=device), + torch.tensor([10, 20], dtype=torch.int64, device=device), + torch.tensor([1.5, 2.5], device=device), ), ) out = tensor.decode() assert isinstance(out, ColumnarTensor) - assert out.shape == (2, 3) + assert out.shape == (3, 3) assert out.device == device - if missing: - rows = out.tolist() - assert rows[0] == [1.5, True, "b"] - assert math.isnan(rows[1][0]) - assert rows[1][1:] == [None, None] - else: - assert out.tolist() == [[1.5, True, "b"], [2.5, False, "a"]] - - -def test_decode_requires_columns() -> None: - tensor = CategoricalTensor( - code=torch.empty(3, 0, dtype=torch.int64), - categories=(), + assert isinstance(out[:, 0], StringTensor) + assert out[:, 0].tolist() == ["a", None, "b"] + assert isinstance(out[:, 1], NullableTensor) + assert out[:, 1].dtype == torch.int64 + assert out[:, 1].tolist() == [20, None, 10] + torch.testing.assert_close( + out[:, 2], + torch.tensor([1.5, torch.nan, 2.5], device=device), + equal_nan=True, ) - with pytest.raises(ValueError, match="at least one column"): - tensor.decode() def test_from_arrow_string_values() -> None: