# Copyright 2024 The MediaPipe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for quantizing tensors.

Note that this is a reduced fork version of the praxis libraries to provide a
self-contained library for packaging.
"""

from typing import Any, List, Optional, Sequence, Tuple, Union

import numpy as np

from mediapipe.tasks.python.genai.converter import external_dependencies

jax = external_dependencies.jax
jnp = external_dependencies.jnp
lax = external_dependencies.lax

JTensor = jax.Array
_UINT4_ZP = 8  # Default zero point for unsigned 4-bit.

# The following values are taken from fake_quantizer.py in Gemax.
_MSE_QUANT_MULS = {
    8: 0.05408,
    4: 0.37755,
}


def _get_scan_range() -> np.ndarray:
  # Produce candidate scan values.
  return np.linspace(1.0, 0.5, num=11)


def _get_mean_error(bound, t, min_value, max_value, p_value):
  scale = bound / max_value
  candidate = jnp.divide(t, scale)
  candidate = jnp.clip(jnp.round(candidate), min_value, max_value)
  candidate = jnp.multiply(candidate, scale)
  pmean_error = jnp.mean(jnp.abs(jnp.subtract(candidate, t)) ** p_value)
  return pmean_error


def _get_best_bound_per_tensor(
    t: JTensor,
    bound: JTensor,
    min_value: float,
    max_value: float,
    p_value: float = 1.0,
) -> JTensor:
  """Scan around [0.5, 1] * hard max value to get bound value for whole tensor.

  This does a scan to get bound value(s) that minimize mean absolute error (MAE)
  between original tensor 't' and quantized tensor. It's (almost) equivalent to
  maximizing entropy.

  Args:
    t: The input float tensor.
    bound: The hard max value for tensor 't'. It has the same length as shape.
    min_value: Minimal value for the quantization bound.
    max_value: Maximal value for the quantization bound.
    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.

  Returns:
    The best bound values for 't', that minimize p-mean error.
  """

  def _quant(scaling_factors):
    return _get_mean_error(
        bound * scaling_factors, t, min_value, max_value, p_value
    )

  scaling_factors = _get_scan_range()
  diffs = jax.vmap(_quant)(scaling_factors)
  best_scaling = scaling_factors[jnp.argmin(diffs)].astype(bound.dtype)
  return bound * best_scaling


def _quantrow(
    vec: JTensor,
    bound: JTensor,
    min_value: float,
    max_value: float,
    p_value: float,
    factors: np.ndarray,
) -> JTensor:
  """Get best rescaling factor from a list of factors applied a channel.

  Args:
    vec: The vector in a channel.
    bound: The hard bound (max(abs(vec))) of the vector.
    min_value: The target min value.
    max_value: The target max value.
    p_value: Exponent of the p-mean error metric.
    factors: The values to be applied on top of bound.

  Returns:
    adjusted bound value out of the list of factors applied to bound.
  """

  def _quant(bounds):
    return _get_mean_error(bounds, vec, min_value, max_value, p_value)

  diffs = jax.vmap(_quant)(bound * factors)
  best_scaling = factors[jnp.argmin(diffs)]
  return bound * best_scaling


def _get_best_bound_per_channel(
    t: JTensor,
    bound: JTensor,
    min_value: float,
    max_value: float,
    p_value: float = 1.0,
) -> JTensor:
  """Scan around [0.5, 1] * hard max value to get bound value for each channel.

  This does a scan to get bound value(s) that minimize mean absolute error (MAE)
  between original tensor 't' and quantized tensor. It's (almost) equivalent to
  maximizing entropy.

  Args:
    t: The input float tensor.
    bound: The hard max value for tensor 't'. It has the same length as shape.
    min_value: Minimal value for the quantization bound.
    max_value: Maximal value for the quantization bound.
    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.

  Returns:
    The best bound values for 't', that minimize p-mean error.
  """
  assert len(t.shape) == 2
  assert len(bound.shape) == 2
  assert t.shape[1] == bound.shape[1]
  assert bound.shape[0] == 1
  scans = _get_scan_range()

  def _quant(tensor, bound, min_value, max_value, p_value, factors):
    ret = np.zeros(bound.shape)
    for i in range(len(tensor)):
      best = _quantrow(
          tensor[i], bound[i], min_value, max_value, p_value, factors
      )
      ret[i] = best
    return ret

  t = t.transpose()
  t_split = list(t)
  res = _quant(t_split, bound[0, :], min_value, max_value, p_value, scans)
  res = res.reshape(bound.shape)
  return res


def get_best_bound(
    t: JTensor,
    bound: JTensor,
    min_value: float,
    max_value: float,
    p_value: float = 1.0,
    per_channel: bool = False,
) -> JTensor:
  """Scan multiple factors on max value to get best bound value.

  This does a scan to get bound value(s) that minimize mean absolute error (MAE)
  between original tensor 't' and quantized tensor. It's (almost) equivalent to
  maximizing entropy.

  Args:
    t: The input float tensor.
    bound: The hard max value for tensor 't'. It has the same length as shape.
    min_value: Minimal value for the quantization bound.
    max_value: Maximal value for the quantization bound.
    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
    per_channel: if get best bound for entire tensor or per channel.

  Returns:
    The best bound values for 't', that minimize p-mean error.
  """
  if per_channel:
    return _get_best_bound_per_channel(t, bound, min_value, max_value, p_value)
  else:
    return _get_best_bound_per_tensor(t, bound, min_value, max_value, p_value)


def get_min_max(
    bits: int = 8,
    unsigned: bool = False,
    use_fp: bool = False,
) -> Tuple[float, float]:
  """Gets the min/max range for a given number of bits.

  Args:
    bits: Target number of bits for quantization.
    unsigned: If True compute min and max for unsigned number, else for signed.
    use_fp: in floating point.

  Returns:
    min/max values for the provide number of bits.
  """
  if use_fp:
    # TODO: support other fp types.
    return -448.0, 448.0
  # Calculation instead of jax.iinfo is used to support bits beside 4 and 8.
  if unsigned:
    # For unsigned 8 bits precision it is [0, 255]
    return 0, 2**bits - 1
  else:
    # For signed 8 bits precision it is [-128, 127]
    return -1 * 2 ** (bits - 1), 2 ** (bits - 1) - 1


def pass_through(x: JTensor, fn: Any) -> JTensor:
  # Create an exactly-zero expression with Sterbenz lemma that has an
  # exactly-one gradient.
  return x - jax.lax.stop_gradient(x) + jax.lax.stop_gradient(fn(x))


def mse_reduce_precision(
    t: JTensor,
    contract_dims: Optional[Sequence[int]],
    need_gradient: bool = False,
    bits: int = 8,
    use_symmetric: bool = True,
    use_fp: bool = False,
) -> Tuple[JTensor, JTensor, Optional[JTensor]]:
  """Reduce the precision of a tensor using MSE quantization.

  Args:
    t: Input tensor.
    contract_dims: Specifies contracting dimensions of the input tensor.
    need_gradient: If gradient is needed out of this function.
    bits: Target number of bits.
    use_symmetric: If the input tensor is quantized symmetrically.
    use_fp: Use floating point.

  Returns:
    A tuple of quantized tensor, quantization scale
      and quantization zero point (optional).
  """
  if need_gradient:
    raise ValueError('need_gradient is not supported for mse_reduce_precision.')
  if use_fp:
    raise ValueError('use_fp is not supported for mse_reduce_precision.')
  if not use_symmetric:
    raise ValueError('Gemax mse_quant is typically symmetric.')

  scale_multiplier = _MSE_QUANT_MULS.get(bits)
  if scale_multiplier is None:
    raise ValueError(f'MSE_QUANT_MULS not defined for {bits} bits.')
  rms = jnp.sqrt(jnp.mean(t**2, axis=contract_dims, keepdims=True))
  scale = scale_multiplier * rms + 1e-12  # Add epsilon as in Gemax

  # Calculate the zero point used in fake_quantizer.mse_quant with narrow range.
  maxq = 2**bits - 1
  zero_mse_quant = (maxq - 1) / 2.0  # This is 7 for 4 bits

  # Simulate the quantization steps from fake_quantizer.mse_quant:
  # 1. Scale and shift: t / scale + zero_mse_quant
  t_shifted = jnp.divide(t, scale) + zero_mse_quant
  # 2. Rounding
  q_rounded = jnp.round(t_shifted)
  # 3. Clipping to [0, maxq - 1]
  q_clipped = jnp.clip(q_rounded, 0, maxq - 1)  # [0, 14] for 4 bits

  # The quantized value (qvar) to be stored should be
  # `q_clipped - zero_mse_quant`. This makes the dequantization match:
  #   (qvar - zp) * scale == (q_clipped - zero_mse_quant) * scale
  qvar = q_clipped - zero_mse_quant  # Ranges from [-7, 7] for 4 bits

  # ml_drift expects a zero point for dequantization. To match the effective
  # dequantization of mse_quant, we set the zero point to 0 (or None).
  zp = None

  container_dtype = (
      jnp.int8 if bits <= 8 else jnp.int16 if bits <= 16 else jnp.int32
  )
  qvar = qvar.astype(container_dtype)

  return qvar, scale, zp


def reduce_precision(
    t: JTensor,
    contract_dims: Optional[Sequence[int]],
    need_gradient: bool = False,
    bits: int = 8,
    optimization_on_bound: bool = False,
    p_value: float = 1.0,
    percentile: float = 1.0,
    use_symmetric: bool = True,
    use_fp: bool = False,
    add_scale_eps: bool = False,
    per_channel: bool = False,
    random_rounding: bool = False,
    key: Optional[jax.Array] = None,
) -> Tuple[JTensor, JTensor, Optional[JTensor]]:
  """Reduce the precision of a tensor.

  Generic for all tensors.

  Args:
    t: Input tensor.
    contract_dims: Specifies contracting dimensions of the input tensor.
    need_gradient: If gradient is needed out of this function.
    bits: Target number of bits.
    optimization_on_bound: If MAE bound optimizer is used.
    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
    percentile: Percentile Factor to apply on the min/max range. Setting this to
      other than 1.0 disables optimization_on_bound.
    use_symmetric: If the input tensor is quantized symmetrically.
    use_fp: Use floating point.
    add_scale_eps: Add eps value or replace zero value by 1 to avoid division by
      zero.
    per_channel: use per-channel clipping optimization.
    random_rounding: round with uniform random.
    key: rng key for rounding.

  Returns:
    A tuple of quantized tensor, quantization scale
      and quantization zero point (optional).
  """
  min_value, max_value = get_min_max(bits, use_fp=use_fp)

  if use_symmetric:
    bound = jnp.max(jnp.abs(t), axis=contract_dims, keepdims=True)
    scale_bound = max_value
  else:
    t_max = jnp.max(t, axis=contract_dims, keepdims=True)
    t_min = jnp.min(t, axis=contract_dims, keepdims=True)
    bound = t_max - t_min
    scale_bound = max_value - min_value

  if percentile < 1.0:
    bound = jnp.multiply(bound, percentile)
  elif optimization_on_bound:
    bound = get_best_bound(
        t, bound, min_value, max_value, p_value, per_channel=per_channel
    )

  scale = bound / scale_bound

  if add_scale_eps:
    # Add epsilon to avoid divide-by-zero.
    scale = scale + jnp.finfo(t.dtype).eps
  else:
    scale = jnp.where(scale == 0.0, 1.0, scale)

  if use_symmetric:
    zp = None
    t = jnp.divide(t, scale)
  else:
    zp = min_value - t_min / scale
    t = jnp.divide(t, scale) + zp
    zp = jnp.multiply(scale, zp)

  if use_fp:
    # No need to round.
    t = jnp.clip(t, min_value, max_value).astype(jnp.float8_e4m3fn)
    # TODO: refactor to remove this logic.
    t = jax.lax.bitcast_convert_type(t, new_dtype=jnp.int8)
  else:
    if need_gradient:
      t = pass_through(t, jnp.round)
      t = jnp.clip(t, min_value, max_value)
    else:
      if random_rounding:
        t = t + jax.random.uniform(
            key=key, shape=t.shape, minval=-0.5, maxval=0.5
        )
      t = jnp.round(t)
      container_dtype = (
          jnp.int8 if bits <= 8 else jnp.int16 if bits <= 16 else jnp.int32
      )
      t = jnp.clip(t, min_value, max_value).astype(container_dtype)

  return t, scale, zp


def quantize_tensor(
    var: np.ndarray,
    axis: List[int],
    factor: float = 1.0,
    sym: bool = True,
    number_bits: int = 8,
    use_fp: bool = False,
    add_scale_eps: bool = False,
    optimization_on_bound: bool = False,
    p_value: float = 1.0,
    per_channel: bool = False,
    block_size: int = 0,
    use_mse_quant: bool = False,
) -> Union[
    Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]
]:
  """Quantize a tensor.

  Args:
    var: The variable to be quantized.
    axis: The axis along which variable will be quantized.
    factor: The clipping factor.
    sym: Symmetric or asymmetric quantize the variable.
    number_bits: Number of bits for quantized value.
    use_fp: do fp with number of bits (i.e. fp8)
    add_scale_eps: add epsilon to scale to avoid division by zero, else it will
      replace zero scale by 1.
    optimization_on_bound: If p-mean bound optimizer is used.
    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
    per_channel: use per-channel clipping optimization.
    block_size:  block size for sub-channel quantization. Defaults to 0, which
      means off.
    use_mse_quant: Whether to use MSE quantization.

  Returns:
    Quantized tensors, along with scales and zero point.
  """
  # TODO: support jnp.float8_e5m2
  assert number_bits == 8 or number_bits == 4
  jnp_var = jnp.asarray(var)
  # When using sub-channel, the contracting dim is split into a sub-channel
  # dim followed by the block dim. Therefore the contracting dim
  # (quantize_axis) should increment by one, and the corresponding pack_dim
  # should also increment by one.
  if block_size > 0:
    # We hard-code in a special case for [0, 1]. It has not been thoroughly
    # tested or verified yet, so we log a warning.
    if len(axis) == 2 and axis == [0, 1]:
      print('Using experimental blockwise int4 axis=[0, 1] compression')
      jnp_var = jnp.reshape(jnp_var, [-1, jnp_var.shape[-1]])
      axis = [0]
    shape = list(jnp_var.shape)
    assert len(axis) == 1, 'Only support 1D sub-channel quantization'
    sub_channels, rem = divmod(shape[axis[0]], block_size)
    assert rem == 0
    shape.insert(axis[0], sub_channels)
    axis[0] += 1
    shape[axis[0]] = block_size
    jnp_var = jnp.reshape(jnp_var, shape)

  if use_mse_quant:
    qvar, scale, zp = mse_reduce_precision(
        jnp_var,
        contract_dims=axis,
        need_gradient=False,
        bits=number_bits,
        use_symmetric=sym,
        use_fp=False,
    )
  else:
    qvar, scale, zp = reduce_precision(
        jnp_var,
        contract_dims=axis,
        need_gradient=False,
        bits=number_bits,
        optimization_on_bound=optimization_on_bound,
        percentile=factor,
        use_symmetric=sym,
        use_fp=use_fp,
        add_scale_eps=add_scale_eps,
        p_value=p_value,
        per_channel=per_channel,
    )
  if sym:
    return np.array(qvar), np.array(jnp.squeeze(scale, axis=axis))  # pytype: disable=wrong-arg-types  # jnp-type
  else:
    return (
        np.array(qvar),
        # CAVEAT: the following squeezes should squeeze along the quantization
        # axis only.
        np.array(jnp.squeeze(scale)),
        np.array(jnp.squeeze(zp)),
    )


def pack_4bit(
    x: np.ndarray, pack_dim: int, packed_dtype: jnp.dtype = jnp.int32
) -> np.ndarray:
  """Pack int8 or uint8 tensor where its values are actually int4 or uint4, to int32 or int8 nibble format along pack_dim.

  Args:
    x: Original int8 or uint8 tensor to pack.
    pack_dim: Dimension to pack along. x.shape[pack_dim] must be divisible by 8,
      when packed_dtype is int32 and divisible by 2 when target_type is int8.
      Also pack_dim must be < x.ndim - 1.
    packed_dtype: Target type to pack to, int32 or int8.

  Returns:
    int32 or int8 packed tensor where the pack_dim size is dividend by 8
    from the original tensor x.
  """
  x = jnp.asarray(x)
  if packed_dtype == jnp.int8 and x.dtype == jnp.uint8:
    # It doesn't make sense to pack uint8 numbers into int4 as we'll
    # the range overlap between uint8 and int4 is [0..7].
    raise ValueError(
        'only int8 input dtype is supported when packing into int8. '
        f'Given {x.dtype}'
    )

  if x.dtype != jnp.int8 and x.dtype != jnp.uint8:
    raise ValueError(
        f'input dtype must be either int8 or uint8. Given {x.dtype}'
    )
  if pack_dim >= x.ndim - 1:
    raise ValueError(
        f'pack_dim must be < input ndim - 1. input shape {x.shape} and pack_dim'
        f' {pack_dim}'
    )
  if packed_dtype != jnp.int32 and packed_dtype != jnp.int8:
    raise ValueError(
        f'packed_dtype must be either int32 or int8. Given {packed_dtype}'
    )
  if packed_dtype == jnp.int32 and x.shape[pack_dim] % 8 != 0:
    raise ValueError(
        'input shape[pack_dim] must be divisible by 8 when target_type '
        f'is int32. Given shape {x.shape}'
    )
  if packed_dtype == jnp.int8 and x.shape[pack_dim] % 2 != 0:
    raise ValueError(
        'input shape[pack_dim] must be divisible by 2 when target_type '
        f'is int8. Given shape {x.shape}'
    )

  int4s_per_packed_type = 8 if packed_dtype == jnp.int32 else 2

  rep_shape = list(x.shape)
  rep_shape.insert(pack_dim + 1, int4s_per_packed_type)
  rep_shape[pack_dim] //= int4s_per_packed_type

  shifts = lax.broadcasted_iota(packed_dtype, rep_shape, pack_dim + 1)
  shifts <<= 2

  # Promote x to packed_dtype
  x = x & jnp.array(0x0F, packed_dtype)
  x = lax.reshape(x, rep_shape)
  x = x << shifts
  x = lax.reduce(x, jnp.array(0x0, packed_dtype), lax.add, [pack_dim + 1])
  return np.asarray(x)


def pack_2bit_into_int8(x: np.ndarray, pack_dim: int) -> np.ndarray:
  """Pack (u)int8 tensor with int2 or uint2 values into int8 along pack_dim.

  Args:
    x: Original int8 or uint8 tensor to pack.
    pack_dim: Dimension to pack along. x.shape[pack_dim] must be divisible by 4.
      Also pack_dim must be < x.ndim - 1.

  Returns:
    int8 packed tensor where the pack_dim size is divided by 4 from the original
    tensor x.
  """
  x = jnp.asarray(x)
  packed_dtype = jnp.int8
  bits = 2
  bits_per_packed_type = 8 // bits  # 4 int2 values per int8

  if x.dtype != jnp.int8 and x.dtype != jnp.uint8:
    raise ValueError(
        f'input dtype must be either int8 or uint8. Given {x.dtype}'
    )
  if pack_dim >= x.ndim - 1:
    raise ValueError(
        f'pack_dim must be < input ndim - 1. input shape {x.shape} and pack_dim'
        f' {pack_dim}'
    )
  if x.shape[pack_dim] % bits_per_packed_type != 0:
    raise ValueError(
        f'input shape[pack_dim] must be divisible by {bits_per_packed_type}'
        f' when packing {bits}-bit values into {packed_dtype}. Given shape'
        f' {x.shape}'
    )

  rep_shape = list(x.shape)
  rep_shape.insert(pack_dim + 1, bits_per_packed_type)
  rep_shape[pack_dim] //= bits_per_packed_type

  # Create shifts: [0, 1, 2, 3]
  shifts = lax.broadcasted_iota(packed_dtype, rep_shape, pack_dim + 1)
  # Adjust shifts for 2-bit packing: [0, 2, 4, 6]
  shifts <<= bits - 1  # Equivalent to shifts * 2

  # Promote x to packed_dtype and mask to keep only the lowest 2 bits (0x03)
  x = x & jnp.array(0x03, packed_dtype)
  x = lax.reshape(x, rep_shape)
  x = x << shifts
  x = lax.reduce(x, jnp.array(0x0, packed_dtype), lax.add, [pack_dim + 1])
  return np.asarray(x)


def update_to_uint4(
    qx: np.ndarray, scale: np.ndarray, zp: Optional[np.ndarray] = None
):
  """Updates the quantized weights from int4 to uint4.

  This is a conversion function designed for XNNPack as it expects the 4-bit
  quantized weight to be represented differently from the original Pax setting.
  Specifically, the differences are:
    1) The dynamic range of weight values: int4 (Pax) vs. uint4 (XNNPack).
    2) The dynamic range of zero-point: float (Pax) vs. uint4 (XNNPack).
    3) The number of zero-point: per-channel (Pax) vs. per-tensor (XNNPack).

  Args:
    qx: np.array of shape [..., channel], which is the quantized weight values
      from Pax in the shape of. The values are in the dynamic range of int4 but
      are hosted as int8 type. Note that if the first dimension is 3, it means
      the qkv matrices are concatenated together and should be treated
      differently.
    scale: np.array of shape [1(3), channel] as np.float type, which are the
      scaling factors for dequantization per channel.
    zp: (optional) np.array of shape [1 (or 3), channel] as np.float type, which
      are the zero points for dequantization per channel.

  Returns:
    A tuple (qx, scale, zp):
      qx: The updated np.array of shape [..., channel] as np.int8 type with
        updated dynamic range as uint4 (with 8 as the default zero points).
      scale: Same as the input scale.
      zp: (optional) np.array of shape [1 (or 3)] as np.int8 type with the
        updated zero point values in the dynamic range as uint4.
  """
  if qx.dtype != np.int8 or ('float' not in str(scale.dtype)):
    raise ValueError(
        'Unexpected dtype qx:' + str(qx.dtype) + ' scale:' + str(scale.dtype)
    )

  scale = scale.astype(np.float32)

  def get_new_zp(old_zp):
    new_zp = old_zp / (scale + np.finfo(np.float32).eps)
    per_tensor_zp = np.mean(new_zp)
    per_tensor_zp = per_tensor_zp.astype(np.int8) + _UINT4_ZP
    return per_tensor_zp

  if zp is not None:
    if qx.shape[0] == 3:
      per_tensor_zp = np.stack([get_new_zp(szp) for szp in zp], axis=0)
    else:
      per_tensor_zp = get_new_zp(zp)
  else:
    per_tensor_zp = (
        _UINT4_ZP * np.ones(shape=(3)) if qx.shape[0] == 3 else _UINT4_ZP
    )

  qx = qx + _UINT4_ZP
  return qx, scale, np.array(per_tensor_zp, dtype=np.int32)


def max_length_axis(shape, axis):
  """Get the axis with the maximum length."""
  if axis is None:
    axis = range(len(shape))
  elif isinstance(axis, int):
    axis = [axis]
  axis = tuple(a % len(shape) for a in axis)
  axis = sorted(axis, reverse=True)  # Prefers larger dim if there is a tie.
  max_len_ind = max(enumerate([shape[a] for a in axis]), key=lambda x: x[1])[0]
  return axis[max_len_ind]


def dequantize_blockwise(values, scales, block_size):
  """Dequantizes a given blockwise-compressed tensor."""
  contracted_dims = []
  if len(values.shape) != len(scales.shape):
    raise ValueError('Values and scales are not the same size')

  # Find the contracted dimensions.
  for i in range(len(values.shape)):
    if values.shape[i] != scales.shape[i]:
      contracted_dims.append(i)

  # Make sure decompressing results in an exact shape match.
  test_shape = list(scales.shape)
  for i in contracted_dims:
    test_shape[i] *= block_size

  if tuple(test_shape) != values.shape:
    raise ValueError('Expanded dimensions did not match')

  # We need to repeat the scale tensor along the target dimension to match the
  # shape of the quantized value tensor.
  target_dim = max_length_axis(values.shape, contracted_dims)
  scales = jnp.repeat(scales, block_size, axis=target_dim)
  return values.astype(jnp.int32).astype(jnp.float32) * scales