gigl.common.utils.feature_quantization.lossless

Lossless narrow-integer packing of feature columns, the lossless half of feature compression.

Many columns of a typical feature set are categorical ids (nn.Embedding indices) that need one to four bytes, not four. FeatureCodec packs a set of such columns into one uint8 row of fixed width. Rows are stored in the packed sidecar that feature quantization already uses (FeatureQuantizationMetadata): GLT’s feature store, RPC, stitch, shared-memory channel and host-to-device copy move opaque byte rows unchanged, and decode rebuilds the float32 columns on the consumer’s device.

Each column has one encoding, which says how the codec stores that column. encode raises on a value its column’s encoding cannot represent:

• float32: exact real values. Also the encoding of a column that preprocessing already quantized.

• bfloat16: real values, lossy. from_stats never picks it.

• int8: ids in [-128, 127].

• uint8: ids in [0, 255].

• uint8_bias1: ids in [-1, 254], stored as id + 1. Fits an out-of-vocabulary id of -1.

• int16: ids in [-32768, 32767].

• int32: ids in [-2**31, 2**31 - 128]. The upper bound is the largest float32 below 2**31.

• bit: exactly 0 or 1 (for example a one-hot column), stored in one bit.

Layout: columns are grouped by encoding and the groups are laid out from widest to narrowest (4-byte, 2-byte, 1-byte) so every group starts at an offset aligned to its element size and can be read with Tensor.view(dtype) on a byte slice. The bit group comes last: its columns are packed eight per byte, column j of the group in bit j % 8 of byte j // 8 (least significant bit first), and the last byte is zero-padded. The row is zero-padded to a multiple of four bytes so that row strides stay aligned for any row subset.

Which columns are stored as narrow integers is inferred from the data: ColumnStats scans the columns and FeatureCodec.from_stats picks the narrowest encoding that holds every value of an integral column. A column that is not integral, does not fit any integer encoding, or holds a -0.0 stays float32 (encode itself writes -0.0 as 0 in an integer encoding). A codec from from_stats is therefore lossless and bit-exact for the data it was built from.

Example

Declare the codec of a node type with five columns: an int16 id, an int8 id, a uint8_bias1 id, an int32 id and a 0/1 bit column.

>>> import torch
>>> from gigl.common.utils.feature_quantization.lossless import FeatureCodec
>>> codec = FeatureCodec(("int16", "int8", "uint8_bias1", "int32", "bit"))
>>> features = torch.tensor(
...     [[300.0, -3.0, -1.0, 100000.0, 1.0], [-2.0, 127.0, 254.0, -7.0, 0.0]]
... )

In the declared column order the columns need 2, 1, 1, 4 bytes and 1 bit. encode lays the groups out widest first: int32 4 bytes, int16 2, int8 1, uint8_bias1 1, then one byte for the bit group. That is 9 bytes, padded to 12, and each row is stored as uint8.

>>> rows = codec.encode(features)
>>> rows.dtype, tuple(rows.shape)
(torch.uint8, (2, 12))

decode returns float32 columns equal to the input.

>>> decoded = codec.decode(rows)
>>> decoded.dtype
torch.float32
>>> bool(torch.equal(decoded, features))
True

A value outside a column’s range is rejected instead of silently wrapping.

>>> codec.encode(torch.tensor([[0.0, 128.0, 0.0, 0.0, 0.0]]))
Traceback (most recent call last):
    ...
ValueError: Column 1 holds a value that is not an integer in [-128, 127] (encoding int8); the packed layout cannot represent it

Classes

ColumnStats	Per-column facts about a float [N, D] matrix that decide which columns can be packed.
FeatureCodec	Packing of a float32 [N, D] matrix into uint8 [N, row_bytes] rows.

Functions

chunk_rows(num_columns, element_size)	Rows per chunk so that one chunk of the input is about _CHUNK_BYTES.
narrowest_integer_encoding(low, high)	The narrowest integer encoding that holds every integer in [low, high], if any.

Module Contents

class gigl.common.utils.feature_quantization.lossless.ColumnStats

Per-column facts about a float [N, D] matrix that decide which columns can be packed.

minimum

[D] column minimums in float64, so a float64 input keeps its exact values (NaN if the column holds a NaN).

maximum

[D] column maximums in float64 (NaN if the column holds a NaN).

integral

[D] bool, True where every value is a whole number (False for fractional values and NaN).

has_negative_zero

[D] bool, True where some value is -0.0.

classmethod empty(width)

Statistics of zero rows, the identity of merge.

Parameters:

width (int)

Return type:

ColumnStats

classmethod from_features(features)

Scans features in row chunks, in its own dtype, so the scan needs no full-size temporaries.

Parameters:

features (jaxtyping.Float[torch.Tensor, rows features_dim])

Return type:

ColumnStats

classmethod merge(stats)

Combines the stats of several partitions (None entries, e.g. empty partitions, are skipped).

Parameters:

stats (Iterable[Optional[ColumnStats]])

Return type:

ColumnStats

has_negative_zero: jaxtyping.Bool[torch.Tensor, width]

integral: jaxtyping.Bool[torch.Tensor, width]

maximum: jaxtyping.Float64[torch.Tensor, width]

minimum: jaxtyping.Float64[torch.Tensor, width]

class gigl.common.utils.feature_quantization.lossless.FeatureCodec

Packing of a float32 [N, D] matrix into uint8 [N, row_bytes] rows.

Parameters:

column_encodings – Encoding of every original column; the module docstring lists the encodings.

decode(rows)

Unpacks uint8 [N, row_bytes] rows (on any device) to float32 [N, D].

Parameters:

rows (jaxtyping.Shaped[torch.Tensor, n _row_bytes])

Return type:

jaxtyping.Float32[torch.Tensor, n {self.num_columns}]

encode(features, *, validate=True)

Packs float32 [N, D] features into uint8 [N, row_bytes] rows.

Parameters:

• features (jaxtyping.Float[torch.Tensor, n _num_columns]) – The [N, D] matrix to pack.

• validate (bool) – Whether to check that every id column holds integers its encoding can store. Pass False only when the statistics this codec was built from cover features, because they already prove that every value fits.

Raises:

ValueError – If the width is not D, if a float32 or bfloat16 column is declared and features is not float32 (the cast would silently lose precision), or if validate is set and an id column holds a non-integral or out-of-range value (including NaN).

Return type:

jaxtyping.UInt8[torch.Tensor, n {self.row_bytes}]

classmethod from_stats(stats)

Picks the narrowest lossless encoding of every column from the data’s statistics.

A column that is integral, holds no -0.0 and fits an integer encoding gets the narrowest such encoding (bit for a column of only 0 and 1). Every other column is float32.

Parameters:

stats (ColumnStats) – Statistics of the data that will be packed, merged over every partition.

Return type:

FeatureCodec

Example

>>> import torch
>>> features = torch.tensor([[0.0, 3.0, 0.5], [1.0, 200.0, 2.0]])
>>> FeatureCodec.from_stats(ColumnStats.from_features(features)).column_encodings
('bit', 'uint8', 'float32')

column_encodings: tuple[str, Ellipsis]

property num_columns: int

Return type:

int

property row_bytes: int

Bytes per packed row, padded up to a multiple of four.

Return type:

int

gigl.common.utils.feature_quantization.lossless.chunk_rows(num_columns, element_size)

Rows per chunk so that one chunk of the input is about _CHUNK_BYTES.

Parameters:

• num_columns (int)

• element_size (int)

Return type:

int

gigl.common.utils.feature_quantization.lossless.narrowest_integer_encoding(low, high)

The narrowest integer encoding that holds every integer in [low, high], if any.

Infinite bounds are held by no encoding.

Parameters:

• low (float)

• high (float)

Return type:

Optional[str]