gigl.common.utils.feature_quantization.torch_ops#

Torch feature dequantization helpers for dataloader collation.

Quantization lives in numpy_ops.py because preprocessing works with CPU arrays in an environment without torch. Dequantization runs in the dataloader collate path, where packed feature data is already represented as torch tensors.

Functions#

dequantize_torch_tensor(packed_features, metadata)

Reconstruct float features from packed uint8 rows.

Module Contents#

gigl.common.utils.feature_quantization.torch_ops.dequantize_torch_tensor(packed_features, metadata)[source]#

Reconstruct float features from packed uint8 rows.

The columns come back in metadata.packed_feature_indices order: first the n-bit quantized columns (approximate), then the lossless columns (exact, except bfloat16 ones).

Parameters:

packed_features (jaxtyping.UInt8[torch.Tensor, ... packed_feature_dim])
metadata (gigl.types.graph.FeatureQuantizationMetadata)

Return type:

jaxtyping.Float32[torch.Tensor, … {len(metadata.packed_feature_indices)}]