Commit 1ddb0530 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

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
parent d2a78f91
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment