%0 = "tf.Less"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
    return %0 : tensor<i1>
  }
}

// -----
// This test verifies that TPUReplicatedInputOps for resource variable args are packed.
module {

      next_use.getOwner()->setOperand(next_use.getOperandNumber(),
                                      private_input.getOutput());
    }
  }

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

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

        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) {