}

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// Checks if `op` is nested in `block`.
bool OpInBlock(Operation* op, Block* block) {

                 replicate_op.GetBody().getArguments())) {
    mlir::TF::TPUReplicatedInputOp input = std::get<0>(input_and_block_arg);
    Value block_arg = std::get<1>(input_and_block_arg);
    mlir::replaceAllUsesInRegionWith(input->getResult(0), block_arg,
                                     cluster.getBody());
    // Update replicated input use in tf.TPUPartitionedInputV2 op.
    for (auto& use : input->getUses()) {

    if (yield->getNumOperands() != block.getNumArguments() + 1)
      return std::nullopt;
    for (auto [yield_operand, block_arg] :
         llvm::zip(yield->getOperands().drop_front(1), block.getArguments())) {
      if (yield_operand != block_arg) return std::nullopt;
    }
  } else {
    // All call results should feed into expected consumer
    // All results of the call should feed into the yield.