94855fb1910ff6b9dd8306dafd097ac9f93eb57b,kfac/python/ops/utils.py,,cross_replica_mean,#Any#Any#,351

Before Change

  
  with tf.name_scope(name, "cross_replica_mean", [tensor]):
    num_shards = get_num_tpu_shards()
    if num_shards is None:
      raise ValueError(
          "Cannot take cross_replica_mean() outside of TPU Context.")
    if num_shards == 1:
      return tensor
    return tf.contrib.tpu.cross_replica_sum(tensor / num_shards)

After Change

  if num_replicas and num_replicas > 1:
    with tf.name_scope(name, "cross_replica_mean", nest.flatten(structure)):
      return nest.map_structure(
          lambda x: tf.contrib.tpu.cross_replica_sum(x / num_replicas),
          structure)
  else:
    return structure

Instances