def sharded_init_partition_spec_fn(
        params, params_partition_spec, partition_spec_for_statistics
    ):
        """Returns a parallel state tree with PartitionSpec associated with state.


        Args:
          params: A pytree with params.
          params_partition_spec: A pytree with PartitionSpec for params.
          partition_spec_for_statistics: PartitionSpec for the statistics.
        """
        # Parallel lists of spec, and params.
        param_pspec_flat, _ = jax.tree_flatten(
            params_partition_spec, is_leaf=lambda x: x is None
        )
        params_flat, treedef = jax.tree_flatten(params)
        assert param_pspec_flat
        assert params_flat
        # Step is replicated across cores.
        # None means cores.
        local_stats_flat = []
        num_statistics = 0
        for param, param_pspec in zip(params_flat, param_pspec_flat):
            param_clone = jnp.zeros(param.shape, dtype=param.dtype)
            preconditioner = preconditioner_from_params(param_clone)
            shapes = preconditioner.shapes_for_preconditioners()
            sizes = []

            index_start = num_statistics
            if not _skip_preconditioning(param):
                sizes = [s[0] for s in shapes]
                shapes = preconditioner.shapes_for_preconditioners()
                num_statistics += len(shapes)

            qdtype = quantized_dtype_for_momentum_buffers(param)
            m1_pspec = param_pspec
            m2_pspec = param_pspec
            m1_scale_pspec = []
            m2_scale_pspec = []
            if qdtype != jnp.float32:
                m1_scale_pspec = _remove_leading_sharding_annotation(m1_pspec)
                m2_scale_pspec = _remove_leading_sharding_annotation(m2_pspec)

            local_stats_flat.append(
                LocalShardedParameterStats(
                    QuantizedValue(
                        param_pspec, [], [], jnp.float32, False, list(param.shape)
                    ),
                    QuantizedValue(
                        m1_pspec, [], m1_scale_pspec, qdtype, False, list(param.shape)
                    ),
                    QuantizedValue(
                        m2_pspec, [], m2_scale_pspec, qdtype, False, list(param.shape)
                    ),
                    init_training_metrics_pspec(),
                    index_start,
                    sizes,
                )
            )

        local_stats = jax.tree_unflatten(treedef, local_stats_flat)
        global_stats = GlobalShardedParameterStats(
            partition_spec_for_statistics,
            partition_spec_for_statistics,
            pjit.PartitionSpec(),
        )
        count_pspec = pjit.PartitionSpec()
        return ShampooState(
            count=count_pspec, stats=ShardedShampooStats(global_stats, local_stats)
        )
Returns a parallel state tree with PartitionSpec associated with state.


        Args:
          params: A pytree with params.
          params_partition_spec: A pytree with PartitionSpec for params.
          partition_spec_for_statistics: PartitionSpec for the statistics.
        
tools/train/scalable_shampoo/distributed_shampoo.py