def preconditioned_grad(self, grad, preconditioners):
        """Precondition the gradient.

        Args:
          grad: A gradient tensor to precondition.
          preconditioners: A list of preconditioners to apply.

        Returns:
          A preconditioned gradient.
        """

        reshaped_grad = jnp.reshape(grad, self._transformed_shape)
        partitioned_grads = self._partitioner.partition(reshaped_grad)
        preconditioned_partitioned_grads = []
        for i, g in enumerate(partitioned_grads):
            should_preconditioned_dims = self.should_precondition_dims()
            num_preconditioners = sum(should_preconditioned_dims)
            preconditioners_for_grad = preconditioners[
                i * num_preconditioners : (i + 1) * num_preconditioners
            ]
            precond_g = g
            rank = len(g.shape)
            for j, precondition in enumerate(should_preconditioned_dims):
                if precondition:
                    precond_g = jnp.tensordot(
                        precond_g, preconditioners_for_grad[j], axes=[[0], [0]]
                    )
                else:
                    precond_g = jnp.transpose(precond_g, axes=(*range(1, rank), 0))
            preconditioned_partitioned_grads.append(precond_g)
        merged_grad = self._partitioner.merge_partitions(
            preconditioned_partitioned_grads
        )
        return jnp.reshape(merged_grad, self._original_shape)
Precondition the gradient.

        Args:
          grad: A gradient tensor to precondition.
          preconditioners: A list of preconditioners to apply.

        Returns:
          A preconditioned gradient.
        
tools/train/scalable_shampoo/distributed_shampoo.py