def _compute_preconditioners(states, params, step):
        """Computes preconditioners for given statistics in states.

        Args:
          states: A list of optimizer states.
          params: A list of params.
          step: Current step number

        Returns:
          New optimizer states after computing the preconditioner.
        """
        statistics = []
        num_statistics_per_state = []
        original_shapes = []
        exponents = []
        max_size = 0
        prev_preconditioners = []

        for state, param in zip(states, params):
            num_statistics = len(state.statistics)
            num_statistics_per_state.append(num_statistics)
            original_shapes_for_state = []
            if num_statistics > 0:
                preconditioner = preconditioner_from_params(param)
                for statistic in state.statistics:
                    exponents.append(
                        preconditioner.exponent_for_preconditioner()
                        if exponent_override == 0
                        else exponent_override
                    )
                    original_shapes_for_state.append(statistic.shape)
                    max_size = max(max_size, statistic.shape[0])

                statistics.extend(state.statistics)
                prev_preconditioners.extend(state.preconditioners)
                original_shapes.extend(original_shapes_for_state)

        if not shard_optimizer_states:
            # Quantization is only enabled if batch_axis_name is not set.
            quantized_dtype = quantized_dtype_for_second_moment_statistics_buffers()

            if quantized_dtype == jnp.float32:
                return _pmap_compute_preconditioners(
                    states,
                    step,
                    statistics,
                    num_statistics_per_state,
                    original_shapes,
                    exponents,
                    max_size,
                    prev_preconditioners,
                )
            else:
                return _pmap_quantized_compute_preconditioners(
                    states,
                    step,
                    statistics,
                    num_statistics_per_state,
                    original_shapes,
                    exponents,
                    max_size,
                    prev_preconditioners,
                )

        else:
            return _pjit_compute_preconditioners(
                states,
                step,
                statistics,
                num_statistics_per_state,
                original_shapes,
                exponents,
                max_size,
                prev_preconditioners,
            )
Computes preconditioners for given statistics in states.

        Args:
          states: A list of optimizer states.
          params: A list of params.
          step: Current step number

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