def sharded_update_fn(grads, state, params):
        """Transform the input gradient and update all statistics in sharded mode.

        Args:
          grads: the gradient tensors for the parameters.
          state: a named tuple containing the state of the optimizer
          params: the parameters that should be updated.

        Returns:
          A tuple containing the new parameters and the new optimizer state.
        """
        params_flat, treedef = jax.tree_flatten(params)
        grads_flat = treedef.flatten_up_to(grads)

        global_stats = state.stats.global_stats
        local_stats_flat = treedef.flatten_up_to(state.stats.local_stats)
        stats_flat = [
            _convert_to_parameter_stats(global_stats, local_stat)
            for local_stat in local_stats_flat
        ]
        new_stats_flat = jax.tree_map(
            lambda g, s, p: _compute_stats(g, s, p, state.count),
            grads_flat,
            stats_flat,
            params_flat,
        )

        outputs = jax.tree_map(
            lambda g, s, p: _transform_grad(g, s, p, state.count),
            grads_flat,
            new_stats_flat,
            params_flat,
        )
        updates_flat, new_stats_flat = list(zip(*outputs)) if outputs else ((), ())

        updates = jax.tree_unflatten(treedef, updates_flat)
        # Create new local_stats
        new_local_stats_flat = [
            _convert_from_parameter_stats(new_stat, local_stat)
            for new_stat, local_stat in zip(new_stats_flat, local_stats_flat)
        ]

        max_size = global_stats.statistics.shape[1]
        new_padded_statistics = []
        for stat in new_stats_flat:
            new_padded_statistics.extend(
                [pad_square_matrix(stat, max_size) for stat in stat.statistics]
            )

        # Create global stats
        # TODO(rohananil): Preconditioner is not updated every step, so cost of
        # stack/pad can be obviated away.
        # 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.
        to_pad = -len(new_padded_statistics) % num_devices_for_pjit
        if not new_padded_statistics:
            to_pad = num_devices_for_pjit
            stat_dtype = jnp.float32
        else:
            stat_dtype = new_padded_statistics[0].dtype

        new_padded_statistics.extend(
            [jnp.eye(max_size, dtype=stat_dtype) for _ in range(to_pad)]
        )
        new_stacked_padded_statistics = jnp.stack(new_padded_statistics)
        new_stacked_padded_statistics = pjit.with_sharding_constraint(
            new_stacked_padded_statistics, statistics_partition_spec
        )

        def _internal_inverse_pth_root_all():
            preconditioners, errors = _matrix_inverse_pth_root_pjit(
                new_stacked_padded_statistics,
                global_stats.exponents,
                statistics_partition_spec,
            )
            return preconditioners, errors

        if preconditioning_compute_steps == 1:
            new_preconditioners, errors = _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 = new_stacked_padded_statistics
            n = new_stacked_padded_statistics.shape[0]
            errors_init = jnp.ones([n], jnp.float32) * inverse_failure_threshold
            init_state = [preconditioners_init, errors_init]
            perform_step = state.count % preconditioning_compute_steps == 0
            new_preconditioners, errors = efficient_cond(
                perform_step, _internal_inverse_pth_root_all, init_state
            )

        new_local_stats_flat = _add_error_into_local_stats(
            new_local_stats_flat, errors, inverse_failure_threshold
        )
        new_local_stats = jax.tree_unflatten(treedef, new_local_stats_flat)
        errors = errors.reshape((-1, 1, 1))
        predicate = jnp.logical_or(
            jnp.isnan(errors), errors >= inverse_failure_threshold
        ).astype(new_preconditioners.dtype)
        # TODO(rohananil): Check for numerical instabilities.
        new_conditional_preconditioners = (
            predicate * global_stats.preconditioners
            + (1.0 - predicate) * new_preconditioners
        )
        new_global_stats = GlobalShardedParameterStats(
            new_stacked_padded_statistics,
            new_conditional_preconditioners,
            global_stats.exponents,
        )
        new_shampoo_state = ShampooState(
            count=state.count + 1,
            stats=ShardedShampooStats(new_global_stats, new_local_stats),
        )
        return updates, new_shampoo_state
Transform the input gradient and update all statistics in sharded mode.

        Args:
          grads: the gradient tensors for the parameters.
          state: a named tuple containing the state of the optimizer
          params: the parameters that should be updated.

        Returns:
          A tuple containing the new parameters and the new optimizer state.
        
tools/train/scalable_shampoo/distributed_shampoo.py