def sharded_init_fn(params):
        """Returns optimizer state (for PJIT mode).

        Args:
          params: the parameters that should be updated.
        """
        params_flat, treedef = jax.tree_flatten(params)
        # Find max size to pad to.
        max_size = 0
        for param in params_flat:
            preconditioner = preconditioner_from_params(param)
            if not _skip_preconditioning(param):
                shapes = preconditioner.shapes_for_preconditioners()
                sizes = [s[0] for s in shapes]
                max_size = max(max(sizes), max_size)

        padded_statistics = []
        padded_preconditioners = []
        local_stats_flat = []
        exponents = []
        for param in params_flat:
            preconditioner = preconditioner_from_params(param)
            shapes = preconditioner.shapes_for_preconditioners()
            sizes = []

            statistics = []
            preconditioners = []
            index_start = len(padded_statistics)
            if not _skip_preconditioning(param):
                sizes = [s[0] for s in shapes]
                shapes = preconditioner.shapes_for_preconditioners()
                statistics = [
                    matrix_epsilon * jnp.eye(max_size, dtype=jnp.float32)
                    for s in shapes
                ]
                preconditioners = [jnp.eye(max_size, dtype=jnp.float32) for s in shapes]
                padded_statistics.extend(statistics)
                padded_preconditioners.extend(preconditioners)
                exponent = (
                    preconditioner.exponent_for_preconditioner()
                    if exponent_override == 0
                    else exponent_override
                )
                exponents.extend([exponent] * len(shapes))

            diagonal_statistics = _quantize_diagonal_statistics(jnp.zeros_like(param))
            diagonal_momentum = _quantize_momentum(jnp.zeros_like(param))
            momentum = _quantize_momentum(jnp.zeros_like(param))

            local_stats_flat.append(
                LocalShardedParameterStats(
                    diagonal_statistics,
                    diagonal_momentum,
                    momentum,
                    init_training_metrics(len(sizes)),
                    index_start,
                    sizes,
                )
            )

        local_stats = jax.tree_unflatten(treedef, local_stats_flat)
        to_pad = -len(padded_statistics) % num_devices_for_pjit
        if max_size == 0:
            to_pad = num_devices_for_pjit
            max_size = block_size
            stat_dtype = jnp.float32
        else:
            stat_dtype = padded_statistics[0].dtype
        # Pad the statistics and preconditioner matrices to be a multiple of
        # num devices.
        # TODO(rohananil): Relax to only the size of the mesh axis where the dim
        # is split on.
        padded_statistics.extend(
            [jnp.eye(max_size, dtype=stat_dtype) for _ in range(to_pad)]
        )
        padded_preconditioners.extend(
            [jnp.eye(max_size, dtype=stat_dtype) for _ in range(to_pad)]
        )
        exponents.extend([1 for _ in range(to_pad)])
        global_stats = GlobalShardedParameterStats(
            jnp.stack(padded_statistics),
            jnp.stack(padded_preconditioners),
            jnp.stack(exponents),
        )
        return ShampooState(
            count=jnp.zeros([], jnp.int32),
            stats=ShardedShampooStats(global_stats, local_stats),
        )
Returns optimizer state (for PJIT mode).

        Args:
          params: the parameters that should be updated.
        
tools/train/scalable_shampoo/distributed_shampoo.py