namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {

using llvm::dyn_cast;
using mlir::Attribute;
using mlir::Block;
using mlir::BlockArgument;
using mlir::DenseIntElementsAttr;
using mlir::failure;
using mlir::Location;
using mlir::LogicalResult;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;

  Value input = shape_op.getInput();
  // If ShapeOp operand is replicate tensor block argument, replace with the
  // associated first replica operand.
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(input)) {
    if (block_arg.getOwner() != replicate_block) return;

    shape_op.setOperand(replicate_op.GetReplicaOperandForBlockArgument(
        block_arg, /*replica=*/0));

    return;
  }

  for (OpResult result : llvm::reverse(op.getResults())) {
    if (!result.use_empty()) continue;
    int result_idx = result.getResultNumber();
    BlockArgument cond_arg = cond.getArgument(result_idx);
    BlockArgument body_arg = cond.getArgument(result_idx);
    Operation *body_ret = body.front().getTerminator();
    // We can eliminate a result if its unused and the corresponding argument

    Block& block = func.front();
    Operation* terminator = block.getTerminator();

    // Replace input_arg(tensor<quant_type>) -> tfr.cast
    // with input_arg(tensor<storage_type>) -> tfr.cast
    for (BlockArgument arg : block.getArguments()) {
      Type arg_type = arg.getType();
      if (auto quant_type = arg_type.cast<TensorType>()
                                .getElementType()

                  QuantizedType::getQuantizedElementType(type))) {
        return true;
      }
    }
  }
  return false;
}

}  // namespace

void QuantizationDriver::InitializeArgState(const BlockArgument arg,
                                            const Value arg_value) {
  const auto [cached, inserted] = value_to_state_.try_emplace(arg_value, 0);
  if (!inserted) {
    arg_states_[arg] = cached->second;

    // Set up a use count for all function arguments. We'll use this to
    // determine whether we have taken care of all uses and can remove
    // the arg.
    for (BlockArgument arg : func.getArguments()) {
      auto uses = arg.getUses();
      use_count[arg.getArgNumber()] = std::distance(uses.begin(), uses.end());
      argument_to_index.insert({arg, arg.getArgNumber()});
    }

    }
  }
  return mlir::success();
}

static void FilterOutBlockArgControlDep(
    ValueRange operands, llvm::SmallVectorImpl<Value> &filtered) {
  for (Value value : operands)
    if (!mlir::isa<mlir::BlockArgument>(value)) filtered.push_back(value);
}

// Split the tfg.NextIteration into tf_executor::NextIterationSourceOp and
// tf_executor::NextIterationSinkOp to break the cycle introduced by itself.

    for (auto entry : llvm::enumerate(cluster_func.getOperands())) {
      auto operand = SkipIdentityAndReadVariable(entry.value());
      auto block_arg = mlir::dyn_cast<BlockArgument>(operand);
      if (block_arg && block_arg.getOwner() == &replicate.GetBody()) {
        // Only mirrored args of ReplicateOp can be annotated.
        if (mirrored_replicate_args.count(block_arg.getArgNumber()) == 0) {

  std::vector<Value> queue;
  queue.push_back(start_node);
  visited.insert(start_node);
  while (!queue.empty()) {
    auto node = queue.back();
    queue.pop_back();
    if (auto arg = mlir::dyn_cast_or_null<BlockArgument>(node)) {
      neighbors.push_back(arg.getArgNumber());
      continue;
    }
    if (!node.getDefiningOp()) continue;
    for (Value operand : node.getDefiningOp()->getOperands()) {

         !op->hasAttrOfType<BoolAttr>(kXlaMapOutsideCompilationAttr);
}

// 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`.