Searched refs:assign_tuple_sharding (Results 1 – 3 of 3) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| H A D | xla_sharding.py | 207 assign_tuple_sharding=False, argument
222 assert use_sharding_op and not assign_tuple_sharding
225 if assign_tuple_sharding:
233 elif assign_tuple_sharding or len(tensor.op.outputs) > 1:
308 def replicate(tensor, assign_tuple_sharding=False, use_sharding_op=False): argument
311 assign_tuple_sharding=assign_tuple_sharding,
317 assign_tuple_sharding=False, argument
322 assign_tuple_sharding=assign_tuple_sharding,
328 assign_tuple_sharding=False, argument
343 assign_tuple_sharding=assign_tuple_sharding,
[all …]
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ |
| H A D | tpu_feed.py | 94 return xla_sharding.replicate(tensor, assign_tuple_sharding=True)
96 return xla_sharding.assign_device(tensor, 0, assign_tuple_sharding=True)
102 assign_tuple_sharding=True)
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/ |
| H A D | xla_compiler.cc | 1099 xla::XlaScopedShardingAssignment assign_tuple_sharding( in BuildArguments() local
|