dst_index);
    }
    return absl::OkStatus();
  }

  auto input_arg = mlir::cast<BlockArgument>(src);
  auto input_node_it = args_.find(input_arg);
  TF_RET_CHECK(input_node_it != args_.end())
      << "Use of BlockArgument encounted before def!";
  // For argument, there is only one result output, so the index is always 0.

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

      // the sense that they're values produced inside this op.
      for (Region& region : op->getRegions()) {
        for (Block& block : region.getBlocks()) {
          for (BlockArgument& arg : block.getArguments()) {
            all_internal_results.insert(arg);
          }
        }
      }
    });
  }

  // All operand values in our set not produced as result by some op in our set.

    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.

               SmallVector<StringRef>(str_values.begin(), str_values.end())));
}

// Creates a new argument for `func_op` that accepts a string tensor containing
// the checkpoint file's prefix.
BlockArgument InsertFilePrefixArgument(func::FuncOp func_op,
                                       OpBuilder& builder) {
  const auto filename_op_type = RankedTensorType::get(

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.