Searched refs:SetOpSharding (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_communication.cc | 214 void SetOpSharding(Operation* op, int64_t tpu_core) { in SetOpSharding() function
268 if (tpu_core) SetOpSharding(send, *tpu_core); in CreateSendOp()
292 if (tpu_core) SetOpSharding(recv, *tpu_core); in CreateRecvOp()
296 if (tpu_core) SetOpSharding(get_tuple_element, *tpu_core); in CreateRecvOp()
302 if (tpu_core) SetOpSharding(new_token, *tpu_core); in CreateRecvOp()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 201 LogicalResult SetOpSharding(Operation* op, Attribute attr, llvm::StringRef name, in SetOpSharding() function
262 if (failed(SetOpSharding(op, input_shardings.getValue()[index], in SetMetadataProtoArgs()
289 if (failed(SetOpSharding(op, output_sharding_and_idx.value(), in SetMetadataProtoRetvals()
|