def quantize(cls, fvalue, quantized_dtype, extract_diagonal=False):
        """Returns quantized value and the bucket."""
        if quantized_dtype == jnp.float32:
            return fvalue, [], []
        elif quantized_dtype == jnp.bfloat16:
            return fvalue.astype(jnp.bfloat16), [], []

        float_dtype = fvalue.dtype
        if quantized_dtype == jnp.int8:
            # value -128 is not used.
            num_buckets = jnp.array(127.0, dtype=float_dtype)
        elif quantized_dtype == jnp.int16:
            # value -32768 is not used.
            num_buckets = jnp.array(32767.0, dtype=float_dtype)
        else:
            raise ValueError(f"Quantized dtype {quantized_dtype} not supported.")
        # max value is mapped to num_buckets

        if extract_diagonal and fvalue.ndim != 2:
            raise ValueError(
                f"Input array {fvalue} must be 2D to work with extract_diagonal."
            )

        diagonal_fvalue = []
        if extract_diagonal:
            diagonal_fvalue = jnp.diag(fvalue)
            # Remove the diagonal entries.
            fvalue = fvalue - jnp.diag(diagonal_fvalue)

        # TODO(rohananil): Extend this by making use of information about the blocks
        # SM3 style which will be useful for diagonal statistics
        # We first decide the scale.
        if fvalue.ndim < 1:
            raise ValueError(
                f"Input array {fvalue} must have a strictly positive number of dimensions."
            )

        max_abs = jnp.max(jnp.abs(fvalue), axis=0)
        bucket_size = max_abs / num_buckets
        bs_expanded = bucket_size[jnp.newaxis, Ellipsis]
        # To avoid divide by 0.0
        bs_nonzero = jnp.where(
            bs_expanded > 0.0, bs_expanded, jnp.ones_like(bs_expanded)
        )
        ratio = fvalue / bs_nonzero
        # We use rounding to remove bias.
        quantized = jnp.round(ratio)
        return quantized.astype(quantized_dtype), diagonal_fvalue, bucket_size
Returns quantized value and the bucket.
tools/train/scalable_shampoo/quantization_utils.py