CategoricalTensor#

class sdm.tensor.CategoricalTensor(code: Tensor, categories: Sequence[Tensor])#

Bases: Tensor

A torch.Tensor for categorical columns.

A CategoricalTensor stores categorical indices in code and one category vector per column in categories. Code values are direct indices into the corresponding category vector. Negative indices represent missing values.

import torch
from sdm import CategoricalTensor, StringTensor

tensor = CategoricalTensor(
    code=torch.randint(0, 2, size=(10, 2)),
    categories=(
        StringTensor.from_list(["USA", "GERMANY"]),
        StringTensor.from_list(["enterprise", "startup"]),
    ),
)

Parameters:

code (Tensor) – The categorical indices of shape [..., C].
categories (Sequence[Tensor]) – A tuple of C category vectors.

Return type:

Self

classmethod from_arrow(array: Array | ChunkedArray, *, dtype: dtype | None = None, device: device | str | None = None) → Self#

Create tensor from a pyarrow.Array.

import pyarrow as pa
from sdm import CategoricalTensor

array = pa.array(["foo", None, "bar"])
tensor = CategoricalTensor.from_arrow(array)

assert tensor.code.tolist() == [[0], [-1], [1]]
assert tensor.categories[0].tolist() == ["foo", "bar"]

Parameters:

array (Array | ChunkedArray) – The pyarrow.Array or pyarrow.ChunkedArray.
dtype (dtype | None) – The dtype.
device (device | str | None) – The device.

Return type:

Self

to_arrow(names: Sequence[str] | None = None) → Table#

Convert this tensor to a two-dimensional pyarrow.Table.

Parameters:: names (Sequence[str] | None) – Column names.
Return type:: Table

classmethod from_cudf(ser: cudf.Series, *, dtype: torch.dtype = torch.int32, device: torch.device | str | None = None) → Self#

Create tensor from a categorical cudf.Series.

Parameters:

ser (cudf.Series) – The categorical cudf.Series.
dtype (torch.dtype) – The dtype.
device (torch.device | str | None) – The device.

Return type:

Self

to_cudf(names: Sequence[str] | None = None) → cudf.DataFrame#

Convert this tensor to a two-dimensional cudf.DataFrame.

Parameters:: names (Sequence[str] | None) – Column names.
Return type:: cudf.DataFrame

classmethod from_tensor(tensor: Tensor) → Self#

Create tensor from a numerical torch.Tensor.

Parameters:: tensor (Tensor) – The numerical tensor.
Return type:: Self

property code: Tensor#: Return the categorical code tensor.

property categories: tuple[Tensor, ...]#: Return category vector for each categorical column.