Searched defs:replicate_op (Results 1 – 7 of 7) sorted by relevance
59 tf_device::ReplicateOp replicate_op, in MoveTailWritesAfterReplicate()104 tf_device::ReplicateOp replicate_op) { in GetTailWritesToReplicateInvariantResourceVars()129 getOperation().walk([&](tf_device::ReplicateOp replicate_op) { in runOnOperation()132 for (auto replicate_op : replicate_ops) { in runOnOperation() local
142 tf_executor::IslandOp island_op, tf_device::ReplicateOp replicate_op, in ExpandReplicateIntoReplicas()239 tf_device::ReplicateOp replicate_op) { in CreateIslandsFromReplicate()322 auto replicate_op = in runOnOperation() local
46 void MakeShapeOpInvariant(tf_device::ReplicateOp replicate_op, int num_replicas, in MakeShapeOpInvariant()124 void HoistReplicateInvariantOps(tf_device::ReplicateOp replicate_op) { in HoistReplicateInvariantOps()
91 auto replicate_op = builder.create<TF::TPUReplicatedInputOp>( in ReorderReplicateAndPartitionedInputs() local
95 tf_device::ReplicateOp replicate_op, OpBuilder* builder) { in ColocateCompositeResourceOpsInReplicate()
479 while_op.body().walk([&](tf_device::ReplicateOp replicate_op) { in runOnOperation()
462 auto replicate_op = builder.create<tf_device::ReplicateOp>( in ReplicateCluster() local