Source code for 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.
"""

import torch
from jaxtyping import Float32, UInt8

from gigl.types.graph import FeatureQuantizationMetadata




def dequantize_torch_tensor(
    packed_features: UInt8[torch.Tensor, "... packed_feature_dim"],
    metadata: FeatureQuantizationMetadata,
) -> Float32[torch.Tensor, "... {metadata.quantized_feature_dim}"]:
    """Reconstruct approximate float features from packed uint8 codes."""
    q = metadata

    if packed_features.size(-1) != q.packed_feature_dim:
        raise ValueError(
            f"Expected packed feature dim {q.packed_feature_dim} for "
            f"{q.quantized_feature_dim} {q.bits}-bit features, got "
            f"{packed_features.size(-1)}."
        )

    codes = _unpack_torch_tensor(
        packed_features, dim=q.quantized_feature_dim, bits=q.bits
    ).float()
    if q.bits == 1:
        if q.neg_mean is None or q.pos_mean is None:
            raise ValueError("1-bit dequantization requires pos_mean/neg_mean")
        return torch.where(codes.bool(), q.pos_mean, q.neg_mean)
    else:
        if q.clip_min is None or q.clip_max is None:
            raise ValueError(f"{q.bits}-bit dequantization requires clip_min/clip_max")
        levels = (1 << q.bits) - 1
        return q.clip_min + (codes / levels) * (q.clip_max - q.clip_min)



def _unpack_torch_tensor(
    packed_features: torch.Tensor, *, dim: int, bits: int
) -> torch.Tensor:
    per_byte = 8 // bits
    mask = (1 << bits) - 1
    # Extract high-bits-first codes from each packed byte.
    shifts = bits * torch.arange(
        per_byte - 1, -1, -1, device=packed_features.device, dtype=torch.uint8
    )
    codes = (packed_features.unsqueeze(-1) >> shifts).bitwise_and(mask)
    return codes.reshape(*packed_features.shape[:-1], -1)[..., :dim]