def quantize(
    x: _torch.Tensor,
    scale: _torch.Tensor,
    zero: _torch.Tensor,
    max_q: _torch.Tensor,
    enable_normal_float: bool,
):
    """
    Quantize ``x`` by rounding and clamping the value using specified
    quantization parameters.
    """
    n_bits = _torch.log2(max_q + 1).item()
    if enable_normal_float:
        if n_bits not in _normal_float_palette:
            raise ValueError(f"Normal float format is not supported for {n_bits}.")
        nf_palette = _normal_float_palette[n_bits]
        nf_palette = nf_palette.to(x.device)
        distances = _torch.cdist((x / scale).view(-1, 1), nf_palette.unsqueeze(0).T)
        indices = _torch.min(distances, dim=1).indices
        return scale * nf_palette[indices].view(x.shape)
    else:
        q = _torch.clamp(_torch.round(x / scale) + zero, 0, max_q)
        return scale * (q - zero)

    Quantize ``x`` by rounding and clamping the value using specified
    quantization parameters.
    
coremltools/optimize/torch/layerwise_compression/_quant.py