Both ParameterServer and Mirrored uses this, and itself is complicated enough.
This also fix a issue that you can't strategy.run(tf.function) under CentralStorageStrategy, by applying the same workaround we have in MirroredStrategy.
PiperOrigin-RevId: 304437119
Change-Id: I6a7a67b88e7a5b7217aa9ffe05882d0ef4097896