K-FAC: Cross Replica Mean for TPU
Adds an op for taking the average of a Tensor across all TPU cores, and uses it before updating covariance statistics. This is a no-op if TPUs aren't used. PiperOrigin-RevId: 179620193
Loading
Please sign in to comment