private_input.getOutput(), [&](OpOperand& use) {
          return forward_pass_ops.contains(use.getOwner());
        });
  }

  VLOG(3) << "Cloned " << cloned_inputs << " TPUReplicatedInputOps";

  // Add TPUReplicatedInput/TPUReplicatedOutput pairs along each edge.
  llvm::SetVector<Operation*> new_forward_ops;
  for (Operation* op : forward_pass_ops) {