def _pmap_quantized_compute_preconditioners(
        states,
        step,
        statistics,
        num_statistics_per_state,
        original_shapes,
        exponents,
        max_size,
        prev_preconditioners,
    ):
        """Computes preconditioners for given statistics in states in PMAP mode.

        For quantization, each statistic is represented by three values:
          quantized matrix, diagonal, and bucket sizes, we run inverse pth-roots
          without ever recreating the original matrix in f32.

        Args:
          states: A list of optimizer states.
          step: Current step number
          statistics: A list of statistics for all variables (for every dim)
          num_statistics_per_state: Number of statistis per state to reconstruct
            output states.
          original_shapes: A list of shapes of the statistics.
          exponents: Exponent power to use for inverse-pth roots.
          max_size: Maximum dim of the statistics to pad.
          prev_preconditioners: Previously available preconditioner.

        Returns:
          New optimizer states after computing the preconditioner.
        """
        num_devices = lax.psum(1, batch_axis_name)
        num_statistics = len(statistics)
        quantized_dtype = quantized_dtype_for_second_moment_statistics_buffers()
        # Complexity here is around: shapes needing be statically shaped,
        # our custom quantization type requires a different type of packing.

        # Parallel tensors:
        # quantized [dxd]
        # diagonals [d] f32
        # bucket_sizes [d] f32
        packed_quantized_statistics = [
            pad_square_matrix(stat.quantized, max_size) for stat in statistics
        ]
        packed_quantized_diagonals = [
            pad_vector(stat.diagonal, max_size) for stat in statistics
        ]
        packed_quantized_bucket_sizes = [
            pad_vector(stat.bucket_size, max_size) for stat in statistics
        ]

        to_pad = -num_statistics % num_devices
        padded_eye = jnp.eye(max_size, dtype=jnp.float32)
        quantized_eye = QuantizedValue.from_float_value(
            padded_eye, quantized_dtype, True
        )
        packed_quantized_statistics.extend(
            [quantized_eye.quantized for _ in range(to_pad)]
        )
        packed_quantized_diagonals.extend(
            [quantized_eye.diagonal for _ in range(to_pad)]
        )
        packed_quantized_bucket_sizes.extend(
            [quantized_eye.bucket_size for _ in range(to_pad)]
        )
        exponents.extend([1 for _ in range(to_pad)])

        if not packed_quantized_statistics:
            return states

        all_quantized_statistics = batch(packed_quantized_statistics, num_devices)
        all_quantized_diagonals = batch(packed_quantized_diagonals, num_devices)
        all_quantized_bucket_sizes = batch(packed_quantized_bucket_sizes, num_devices)
        all_exponents = batch(exponents, num_devices)

        def _internal_inverse_pth_root_all():
            current_replica = lax.axis_index(batch_axis_name)
            (
                quantized_preconditioners,
                quantized_diagonals,
                quantized_bucket_sizes,
                errors,
            ) = _quantized_matrix_inverse_pth_root_vmap(
                all_quantized_statistics[current_replica],
                all_quantized_diagonals[current_replica],
                all_quantized_bucket_sizes[current_replica],
                all_exponents[current_replica],
            )
            quantized_preconditioners = jax.lax.all_gather(
                quantized_preconditioners, batch_axis_name
            )
            quantized_diagonals = jax.lax.all_gather(
                quantized_diagonals, batch_axis_name
            )
            quantized_bucket_sizes = jax.lax.all_gather(
                quantized_bucket_sizes, batch_axis_name
            )
            errors = jax.lax.all_gather(errors, batch_axis_name)
            quantized_preconditioners_flat = unbatch(quantized_preconditioners)
            quantized_diagonals_flat = unbatch(quantized_diagonals)
            quantized_bucket_sizes_flat = unbatch(quantized_bucket_sizes)
            errors_flat = unbatch(errors)
            return (
                quantized_preconditioners_flat,
                quantized_diagonals_flat,
                quantized_bucket_sizes_flat,
                errors_flat,
            )

        if preconditioning_compute_steps == 1:
            (
                quantized_preconditioners_flat,
                quantized_diagonals_flat,
                quantized_bucket_sizes_flat,
                errors_flat,
            ) = _internal_inverse_pth_root_all()
        else:
            # Passing statistics instead of preconditioners as they are similarly
            # shaped tensors. Note statistics will be ignored as we are passing in
            # a large init value for error.
            quantized_preconditioners_init = packed_quantized_statistics
            quantized_diagonals_init = packed_quantized_diagonals
            quantized_bucket_sizes_init = packed_quantized_bucket_sizes
            errors_init = [inverse_failure_threshold] * len(
                quantized_preconditioners_init
            )
            init_state = [
                quantized_preconditioners_init,
                quantized_diagonals_init,
                quantized_bucket_sizes_init,
                errors_init,
            ]
            perform_step = step % preconditioning_compute_steps == 0
            (
                quantized_preconditioners_flat,
                quantized_diagonals_flat,
                quantized_bucket_sizes_flat,
                errors_flat,
            ) = efficient_cond(perform_step, _internal_inverse_pth_root_all, init_state)

        def _skip(error):
            condition = jnp.logical_or(
                jnp.isnan(error), error >= inverse_failure_threshold
            )
            return condition.astype(error.dtype)

        def _select_preconditioner(error, new_p, old_p):
            return lax.cond(
                _skip(error), lambda _: old_p, lambda _: new_p, operand=None
            )

        new_quantized_preconditioners_flat = []
        new_quantized_diagonals_flat = []
        new_quantized_bucket_sizes_flat = []
        new_errors_flat = []
        for p, d, b, shape, prev_p, error in zip(
            quantized_preconditioners_flat,
            quantized_diagonals_flat,
            quantized_bucket_sizes_flat,
            original_shapes,
            prev_preconditioners,
            errors_flat,
        ):
            new_quantized_preconditioners_flat.append(
                _select_preconditioner(
                    error, p[: shape[0], : shape[1]], prev_p.quantized
                )
            )
            new_quantized_diagonals_flat.append(
                _select_preconditioner(error, d[: shape[0]], prev_p.diagonal)
            )
            new_quantized_bucket_sizes_flat.append(
                _select_preconditioner(error, b[: shape[0]], prev_p.bucket_size)
            )
            new_errors_flat.append(error)

        assert len(states) == len(num_statistics_per_state)
        assert len(new_quantized_preconditioners_flat) == num_statistics
        assert len(new_quantized_diagonals_flat) == num_statistics
        assert len(new_quantized_bucket_sizes_flat) == num_statistics

        # Add back empty preconditioners so we that we can set the optimizer state.
        preconditioners_for_states = []
        errors_for_states = []
        idx = 0
        for num_statistics, state in zip(num_statistics_per_state, states):
            if num_statistics == 0:
                preconditioners_for_states.append([])
                errors_for_states.append(jnp.array(0, jnp.float32))
            else:
                quantized_preconditioners_for_state = (
                    new_quantized_preconditioners_flat[idx : idx + num_statistics]
                )
                quantized_diagonals_for_state = new_quantized_diagonals_flat[
                    idx : idx + num_statistics
                ]
                quantized_bucket_sizes_for_state = new_quantized_bucket_sizes_flat[
                    idx : idx + num_statistics
                ]
                errors_for_state = jnp.stack(
                    new_errors_flat[idx : idx + num_statistics]
                )

                assert len(state.statistics) == len(quantized_preconditioners_for_state)
                assert len(state.statistics) == len(quantized_diagonals_for_state)
                assert len(state.statistics) == len(quantized_bucket_sizes_for_state)
                assert len(state.statistics) == len(errors_for_state)

                quantized_preconditioners = []
                for qv, qd, qb in zip(
                    quantized_preconditioners_for_state,
                    quantized_diagonals_for_state,
                    quantized_bucket_sizes_for_state,
                ):
                    quantized_preconditioners.append(
                        QuantizedValue(qv, qd, qb, qv.dtype, True, list(qv.shape))
                    )
                preconditioners_for_states.append(quantized_preconditioners)
                errors_for_states.append(errors_for_state)
                idx += num_statistics
        new_states = []
        for state, new_preconditioners, new_errors in zip(
            states, preconditioners_for_states, errors_for_states
        ):
            if state.statistics:
                new_errors = jnp.where(
                    jnp.logical_and(
                        new_errors > 0.0, new_errors != inverse_failure_threshold
                    ),
                    new_errors,
                    state.training_metrics.inverse_pth_root_errors,
                )
            new_training_metrics = TrainingMetrics(new_errors)
            new_states.append(
                ParameterStats(
                    state.diagonal_statistics,
                    state.statistics,
                    new_preconditioners,
                    state.diagonal_momentum,
                    state.momentum,
                    new_training_metrics,
                )
            )

        return new_states
Computes preconditioners for given statistics in states in PMAP mode.

        For quantization, each statistic is represented by three values:
          quantized matrix, diagonal, and bucket sizes, we run inverse pth-roots
          without ever recreating the original matrix in f32.

        Args:
          states: A list of optimizer states.
          step: Current step number
          statistics: A list of statistics for all variables (for every dim)
          num_statistics_per_state: Number of statistis per state to reconstruct
            output states.
          original_shapes: A list of shapes of the statistics.
          exponents: Exponent power to use for inverse-pth roots.
          max_size: Maximum dim of the statistics to pad.
          prev_preconditioners: Previously available preconditioner.

        Returns:
          New optimizer states after computing the preconditioner.
        
tools/train/scalable_shampoo/distributed_shampoo.py