def updated_statistics_from_grad(
        self,
        stats,
        grad,
        w1,
        w2,
        to_float=None,
        from_float=None,
        precision=None,
    ):
        """Update statistics from gradients.

        Args:
          stats: Old statistics or its Cholesky factor if `cholesky` is True.
          grad: Gradient to compute statistics from.
          w1: Weight for old statistics.
          w2: Weight for new statistics.
          to_float: Optional function for converting stats to floating point.
          from_float: Optional function for converting from floating point.
          precision: Optional precision XLA related flag, the available options are:
            a) lax.Precision.DEFAULT (better step time, but not precise)
            b) lax.Precision.HIGH (increased precision, slower)
            c) lax.Precision.HIGHEST (best possible precision, slowest)

        Returns:
          A list of updated gradient statistics for each partition.
        """
        to_float = to_float if to_float is not None else (lambda x: x)
        from_float = from_float if from_float is not None else (lambda x: x)
        update = functools.partial(gram_weighted_update, precision=precision)
        reshaped_grad = jnp.reshape(grad, self._transformed_shape)
        partitioned_grads = self._partitioner.partition(reshaped_grad)
        new_stats = []
        index = 0
        for g in partitioned_grads:
            should_preconditioned_dims = self.should_precondition_dims()
            num_preconditioners = sum(should_preconditioned_dims)
            for axis in range(num_preconditioners):
                new_stat = update(to_float(stats[index]), g, axis, w1, w2)
                new_stats.append(from_float(new_stat))
                index += 1
        return new_stats
Update statistics from gradients.

        Args:
          stats: Old statistics or its Cholesky factor if `cholesky` is True.
          grad: Gradient to compute statistics from.
          w1: Weight for old statistics.
          w2: Weight for new statistics.
          to_float: Optional function for converting stats to floating point.
          from_float: Optional function for converting from floating point.
          precision: Optional precision XLA related flag, the available options are:
            a) lax.Precision.DEFAULT (better step time, but not precise)
            b) lax.Precision.HIGH (increased precision, slower)
            c) lax.Precision.HIGHEST (best possible precision, slowest)

        Returns:
          A list of updated gradient statistics for each partition.
        
tools/train/scalable_shampoo/distributed_shampoo.py