Searched defs:replicate_op (Results 1 – 8 of 8) sorted by relevance
140 tf_executor::IslandOp island_op, tf_device::ReplicateOp replicate_op, in ExpandReplicateIntoReplicas()237 tf_device::ReplicateOp replicate_op) { in CreateIslandsFromReplicate()320 auto replicate_op = in runOnFunction() local
76 void MakeShapeOpInvariant(tf_device::ReplicateOp replicate_op, int num_replicas, in MakeShapeOpInvariant()154 void HoistReplicateInvariantOps(tf_device::ReplicateOp replicate_op) { in HoistReplicateInvariantOps()
42 module.walk([&](tf_device::ReplicateOp replicate_op) { in runOnOperation()
95 auto replicate_op = builder.create<TF::TPUReplicatedInputOp>( in ReorderReplicateAndPartitionedInputs() local
92 tf_device::ReplicateOp replicate_op, OpBuilder* builder) { in ColocateCompositeResourceOpsInReplicate()
80 getOperation().walk([&](TF::TPUReplicateMetadataOp replicate_op) { in runOnOperation()
526 while_op.body().walk([&](tf_device::ReplicateOp replicate_op) { in runOnOperation()
424 auto replicate_op = builder.create<tf_device::ReplicateOp>( in ReplicateCluster() local