def sharded_init_shape_and_dtype_fn(params):
        """Returns a parallel state tree with shape, dtype associated with state.


        Args:
          params: A pytree with params.
        """
        # Parallel lists of spec, and params.
        params_flat, treedef = jax.tree_flatten(params)
        assert params_flat
        # Step is replicated across cores.
        # None means cores.
        local_stats_flat = []
        num_statistics = 0
        for param in params_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_shape_and_dtype = [list(param.shape), param.dtype]
            m2_shape_and_dtype = [list(param.shape), param.dtype]
            m1_scale_shape_and_dtype = []
            m2_scale_shape_and_dtype = []
            if qdtype != jnp.float32:
                m1_scale_shape_and_dtype = [list(param.shape)[1:], qdtype]
                m2_scale_shape_and_dtype = [list(param.shape)[1:], qdtype]

            diagonal_statistics_shape_and_dtype = [list(param.shape), param.dtype]
            local_stats_flat.append(
                LocalShardedParameterStats(
                    QuantizedValue(
                        diagonal_statistics_shape_and_dtype,
                        [],
                        [],
                        jnp.float32,
                        False,
                        list(param.shape),
                    ),
                    QuantizedValue(
                        m1_shape_and_dtype,
                        [],
                        m1_scale_shape_and_dtype,
                        qdtype,
                        False,
                        list(param.shape),
                    ),
                    QuantizedValue(
                        m2_shape_and_dtype,
                        [],
                        m2_scale_shape_and_dtype,
                        qdtype,
                        False,
                        list(param.shape),
                    ),
                    init_training_metrics_shapes(len(sizes)),
                    index_start,
                    sizes,
                )
            )

        local_stats = jax.tree_unflatten(treedef, local_stats_flat)
        max_statistics_size = _max_statistics_size_from_params(params_flat)
        to_pad = -num_statistics % num_devices_for_pjit
        num_statistics += to_pad
        if num_statistics == 0:
            num_statistics = num_devices_for_pjit
            max_statistics_size = block_size
        statistics_shape = [num_statistics, max_statistics_size, max_statistics_size]
        global_stats = GlobalShardedParameterStats(
            [statistics_shape, jnp.float32],
            [statistics_shape, jnp.float32],
            [[num_statistics], jnp.int32],
        )
        return ShampooState(
            count=[[], jnp.float32],
            stats=ShardedShampooStats(global_stats, local_stats),
        )
Returns a parallel state tree with shape, dtype associated with state.


        Args:
          params: A pytree with params.
        
tools/train/scalable_shampoo/distributed_shampoo.py