def update_fn(grads, state, params):
        """Transform the input gradient and update all statistics.

        Args:
          grads: the gradient tensors for the parameters
            and any custom gradients for preconditioners.
          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)
        stats_flat = treedef.flatten_up_to(state.stats)
        grads_flat = treedef.flatten_up_to(grads)
        stats_grads = grads_flat

        new_stats_flat = jax.tree_map(
            lambda g, s, p: _compute_stats(g, s, p, state.count),
            stats_grads,
            stats_flat,
            params_flat,
        )

        new_stats_flat = _compute_preconditioners(
            new_stats_flat, params_flat, state.count
        )
        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)
        new_stats = jax.tree_unflatten(treedef, new_stats_flat)

        new_state = ShampooState(count=state.count + 1, stats=new_stats)
        return updates, new_state
Transform the input gradient and update all statistics.

        Args:
          grads: the gradient tensors for the parameters
            and any custom gradients for preconditioners.
          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