Searched refs:cross_replica_sum (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 92 def cross_replica_sum(x, group_assignment=None, name=None): function
108 return gen_tpu_ops.cross_replica_sum(x, group_assignment, name=name)
150 return [gen_tpu_ops.cross_replica_sum(grad, op.inputs[1]), None]
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | ar_crs_combiner.h | 91 ArCrsPair(HloInstruction* all_reduce, HloInstruction* cross_replica_sum, in ArCrsPair()
93 : ar(all_reduce), crs(cross_replica_sum), distance(dist) {} in ArCrsPair()
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_optimizer.py | 169 summed_grads_and_vars.append((tpu_ops.cross_replica_sum(
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 449 return tpu_ops.cross_replica_sum(value)
|
| /external/tensorflow/tensorflow/contrib/tpu/python/tpu/ |
| D | keras_support.py | 318 concat_tensor = tpu_ops.cross_replica_sum(local_tensor_with_holes)
352 return [tpu_ops.cross_replica_sum(grad) / num_shards for grad in grads]
|