Function embedding_2bit

exir/passes/_quant_patterns_and_replacements.py:248–285

(
    weight: torch.Tensor,
    weight_scales: torch.Tensor,
    weight_zero_points: Optional[torch.Tensor],
    weight_quant_min: int,
    weight_quant_max: int,
    indices: torch.Tensor,
)

@impl(quantized_decomposed_lib, "embedding_2bit", "CompositeExplicitAutograd")
def embedding_2bit(
    weight: torch.Tensor,
    weight_scales: torch.Tensor,
    weight_zero_points: Optional[torch.Tensor],
    weight_quant_min: int,
    weight_quant_max: int,
    indices: torch.Tensor,
) -> torch.Tensor:
    assert (
        weight_quant_min == -2
    ), "embedding_2bit in ExecuTorch expects weight_quant_min == -2"
    assert (
        weight_quant_max == 1
    ), "embedding_2bit in ExecuTorch expects weight_quant_max == 1"

    embedding_weight_checks(weight, weight_scales, weight_zero_points)
    group_size = (4 * weight.size(1)) // (
        weight_scales.size(1) if weight_scales.dim() == 2 else 1
    )
    weight_0 = weight & 3
    weight_1 = (weight & 12) >> 2
    weight_2 = (weight & 48) >> 4
    weight_3 = (weight & 192) >> 6
    weight_unpacked = torch.stack((weight_0, weight_1, weight_2, weight_3), dim=-1)
    weight = weight_unpacked.view(weight.shape[0], -1)
    weight = weight.view(torch.int8).add(-2)

    weight = torch.ops.quantized_decomposed.dequantize_per_channel_group.default(
        weight,
        weight_scales,
        weight_zero_points,
        weight_quant_min,
        weight_quant_max,
        weight.dtype,
        group_size,
        weight_scales.dtype,
    )
    return torch.ops.aten.embedding.default(weight, indices)


@register_fake("quantized_decomposed::embedding_2bit")