include "mlir/IR/CommonTypeConstraints.td"

def CloneF32ElementsAttrWithOnes
  : NativeCodeCall<"DenseElementsAttr::get($0.getType().cast<ShapedType>(), (float)1.0)">;

def NotConstant : Constraint<
    CPred<"$0.isa<BlockArgument>() || !llvm::isa<stablehlo::ConstantOp>($0.getDefiningOp())">,
    "Is not a constant.">;

def : Pat<(StableHLO_DivOp $l,
            (StableHLO_ConstantOp:$divisor FloatElementsAttr<32>:$cst)),

    llvm::SetVector<Operation*> ops_to_erase;
    llvm::MapVector<BlockArgument, TF::TPUReplicatedInputOp> cache;
    for (auto input_op : func.getRegion().getOps<TF::TPUReplicatedInputOp>()) {
      // We're only expecting a single input argument to be replicated.
      if (input_op->getNumOperands() > 1) continue;
      Value operand = input_op->getOperand(0);
      if (!llvm::isa<BlockArgument>(operand)) continue;

class Block:
    def __init__(self, *args, **kwargs) -> None: ...
    def addArgument(self, *args, **kwargs) -> Any: ...
    def end(self) -> Block_Iterator: ...
    def new(self) -> Block: ...

class BlockArgument(Value):
    def __init__(self, *args, **kwargs) -> None: ...

class Block_Iterator:
    def __init__(self, *args, **kwargs) -> None: ...

class Builder:
    def __init__(self, arg0: MLIRContext) -> None: ...

                              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

          ExtractTPUCopyWithDynamicShapeOpPass> {
  void runOnOperation() override;
};

// 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();
}

// Check if the TPUCopyWithDynamicShapeOp is valid.

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