def METHOD_NAME(codes: JTensor,
                            num_classes: int,
                            *,
                            paddings: JTensor,
                            data_parallel_axis: Optional[str] = None):
  """Computes codebook coverage within a batch.
  Args:
    codes:         [..., num_groups], values are in [0, num_classes).
    num_classes:   A Python int.
    paddings:      [...], 0/1 value tensor.
    data_parallel_axis: If set will psum() over the axis
  Returns:
    A scalar tf.Tensor, avg coverage across groups.
  """
  # [num_groups, num_classes]
  _, _, histogram = batch_pplx_entropy_from_codes(
      codes,
      num_classes,
      paddings=paddings,
      data_parallel_axis=data_parallel_axis)
  onehot = jnp.greater(histogram, 0).astype(jnp.float32)
  avg_num_covered_words = jnp.mean(jnp.sum(onehot, -1))
  return avg_num_covered_words / num_classes