def _pmap_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.

        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.
        """
        if batch_axis_name:
            num_devices = lax.psum(1, batch_axis_name)
        else:
            num_devices = 1
        num_statistics = len(statistics)
        # Pad statistics and exponents to next multiple of num_devices.
        packed_statistics = [pad_square_matrix(stat, max_size) for stat in statistics]
        to_pad = -num_statistics % num_devices
        packed_statistics.extend(
            [jnp.eye(max_size, dtype=packed_statistics[0].dtype) for _ in range(to_pad)]
        )
        exponents.extend([1 for _ in range(to_pad)])

        if not packed_statistics:
            return states

        all_statistics = batch(packed_statistics, num_devices)
        all_exponents = batch(exponents, num_devices)

        def _internal_inverse_pth_root_all():
            if batch_axis_name:
                current_replica = lax.axis_index(batch_axis_name)
                preconditioners, errors = _matrix_inverse_pth_root_vmap(
                    all_statistics[current_replica], all_exponents[current_replica]
                )
                preconditioners = jax.lax.all_gather(preconditioners, batch_axis_name)
                errors = jax.lax.all_gather(errors, batch_axis_name)
                preconditioners_flat = unbatch(preconditioners)
                errors_flat = unbatch(errors)
            else:
                preconditioners, errors = _matrix_inverse_pth_root_vmap(
                    all_statistics[0], all_exponents[0]
                )
                preconditioners_flat = unbatch(jnp.stack([preconditioners]))
                errors_flat = unbatch(jnp.stack([errors]))

            return preconditioners_flat, errors_flat

        if preconditioning_compute_steps == 1:
            preconditioners_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.
            preconditioners_init = packed_statistics
            errors_init = [inverse_failure_threshold] * len(packed_statistics)
            init_state = [preconditioners_init, errors_init]
            perform_step = step % preconditioning_compute_steps == 0
            preconditioners_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_preconditioners_flat = []
        new_errors_flat = []
        for p, shape, prev_p, error in zip(
            preconditioners_flat, original_shapes, prev_preconditioners, errors_flat
        ):
            new_preconditioners_flat.append(
                _select_preconditioner(error, p[: shape[0], : shape[1]], prev_p)
            )
            new_errors_flat.append(error)

        assert len(states) == len(num_statistics_per_state)
        assert len(new_preconditioners_flat) == num_statistics
        assert len(new_errors_flat) == num_statistics

        # Add back empty preconditioners so we that we can set the optimizer state.
        preconditioners_for_states = []
        idx = 0
        errors_for_states = []
        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:
                preconditioners_for_state = new_preconditioners_flat[
                    idx : idx + num_statistics
                ]
                assert len(state.statistics) == len(preconditioners_for_state)
                preconditioners_for_states.append(preconditioners_for_state)

                errors_for_state = jnp.stack(
                    new_errors_flat[idx : idx + num_statistics]
                )
                assert len(state.statistics) == len(errors_for_state)
                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.

        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