Searched refs:use_sharding_op (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 167 use_sharding_op=False): argument
179 if use_sharding_op:
234 def copy_sharding(from_tensor, to_tensor, use_sharding_op=False): argument
249 if use_sharding_op:
261 def replicate(tensor, assign_tuple_sharding=False, use_sharding_op=False): argument
265 use_sharding_op=use_sharding_op)
271 use_sharding_op=False): argument
276 use_sharding_op=use_sharding_op)
282 use_sharding_op=False): argument
295 use_sharding_op=use_sharding_op)
[all …]
|
| /external/tensorflow/tensorflow/python/training/ |
| D | slot_creator_test.py | 185 v, np.array([0, 1]), [0], use_sharding_op=False)
197 v, np.array([0, 1]), [0], use_sharding_op=False)
|
| D | slot_creator.py | 112 slot = xla_sharding.copy_sharding(primary, slot, use_sharding_op=False)
|
| D | adam_test.py | 395 v, np.array([0, 1]), [0], use_sharding_op=False)
|
| D | moving_averages_test.py | 502 v = xla_sharding.mesh_split(v, np.array([0, 1]), [0], use_sharding_op=False)
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 458 tensor, logical_device_id, use_sharding_op=True)
548 return xla_sharding.tile(tensor, tile_assignment, use_sharding_op=True)
596 return xla_sharding.replicate(tensor, use_sharding_op=True)
|