Region& second_region) {
    if (!mlir::isa<BlockArgument>(first) || !mlir::isa<BlockArgument>(second))
      return false;
    BlockArgument first_block_arg = mlir::cast<BlockArgument>(first);
    BlockArgument second_block_arg = mlir::cast<BlockArgument>(second);

    // 2 block arguments will match if they are the same argument number, and

    }
  } else {
    info.decomposed_callee = lowered_callee;
    for (auto& entry : callee_map) {
      info.stack_var_arg_to_size_arg[mlir::cast<BlockArgument>(entry.getFirst())
                                         .getArgNumber()] =
          mlir::cast<BlockArgument>(entry.getSecond()).getArgNumber();
    }
    if (lowered_callee != callee) {
      // Add the clone with a new name.

namespace {
constexpr char kDeviceAttr[] = "device";

// A type that abstracts over types that have uses accessible via `getUses`.
using Source = PointerUnion<Operation *, BlockArgument *>;

// We use union-find algorithm to build clusters of connected operations based
// on the user provided policy. If an operation can be clustered (one of the

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;

  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

                  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;

    }
  }
  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.

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

using OpShardingVector = llvm::SmallVector<OpShardingVariant, 8>;
using OptionalOpShardingVector =
    llvm::SmallVector<std::optional<OpShardingVariant>, 8>;
using llvm::StringRef;
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::Builder;
using mlir::IntegerAttr;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OpOperand;
using mlir::OpResult;