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

void TPUAnnotateDynamicShapeInputsPass::runOnOperation() {
  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());

int64_t GetTransferredTensorBytes(func::CallOp from_graph,
                                  func::CallOp to_graph) {
  int64_t total_size_transferred = 0;
  for (auto input : to_graph.getOperands()) {
    Operation* input_op = input.getDefiningOp();
    if (input_op && input_op == from_graph.getOperation()) {
      auto input_type =
          mlir::dyn_cast_or_null<RankedTensorType>(input.getType());

    DenseFPElementsAttr min_value, max_value;
    if (auto id1 = dyn_cast_or_null<TF::IdentityOp>(min.getDefiningOp())) {
      id1.replaceAllUsesWith(id1.getInput());
      min = tf_op.getMin();
      rewriter.eraseOp(id1);
    }
    if (auto id2 = dyn_cast_or_null<TF::IdentityOp>(max.getDefiningOp())) {
      id2.replaceAllUsesWith(id2.getInput());
      max = tf_op.getMax();
      rewriter.eraseOp(id2);
    }

    // case is useful for float16 quantization. Since all ops have been
    // legalized to tflite ops, so we only care about ConstOp or QConstOp or
    // mlir constant op.
    Operation* input_op = operand.getDefiningOp();
    if (input_op == nullptr) return failure();

    Attribute attr;
    if (matchPattern(input_op, m_Constant(&attr))) {
      // Constant case.
      builder->setInsertionPoint(op);

                  AttrType &max_value) const {
    Value min = tf_op.getMin(), max = tf_op.getMax();
    if (auto min_id = min.getDefiningOp<TF::IdentityOp>()) {
      min = min_id.getInput();
    }
    if (auto max_id = max.getDefiningOp<TF::IdentityOp>()) {
      max = max_id.getInput();
    }

    if (!matchPattern(min, m_Constant(&min_value))) {
      return false;
    }

#define COMPOSITE_UTILS_TD

include "mlir/IR/PatternBase.td"

// See the function doc in the header file.
def GetNhwcReturnTypeFromNchw: NativeCodeCall<
  "GetNhwcReturnTypeFromNchw((*$0.begin()).getDefiningOp())">;

// When given a DenseIntElementsAttr containing I64 elements, this extracts
// one I32IntegerAttr from the given index.
class GetI32At<int index>: NativeCodeCall<

    auto assign_variable_op =
        dyn_cast_or_null<AssignVariableOp>(var_handle_user);
    if (!assign_variable_op) continue;
    auto value_op = assign_variable_op.getValue().getDefiningOp();
    auto dq_op = dyn_cast_or_null<DequantizeOp>(value_op);
    if (!dq_op || ref_qtype) continue;
    ref_qtype = dq_op.getInput().getType();
  }
  return ref_qtype;
}

class QuantizeVariablesPass

  resource_infos_[resource].potentially_written = true;
  auto* operation = resource.getDefiningOp();
  if (operation && isa<TF::VarHandleOp>(operation)) {
    mutable_variables_.insert(GetResourceKey(operation));
  }
}

bool ResourceAnalyzer::IsPotentiallyWritten(Value resource) const {
  assert(IsResource(resource));
  auto* operation = resource.getDefiningOp();
  if (operation && isa<TF::VarHandleOp>(operation))

    Value pad_input = pad_op.getOperand();
    RankedTensorType pad_type = pad_op.getType().cast<RankedTensorType>();

    auto transpose_op = pad_input.getDefiningOp<stablehlo::TransposeOp>();
    if (!transpose_op || !transpose_op->hasOneUse()) return failure();
    Value transpose_input = transpose_op.getOperand();

    ArrayRef<int64_t> transpose_perm = transpose_op.getPermutation();

  LogicalResult match(mlir::stablehlo::TransposeOp op) const override {
    Value operand = op.getOperand();
    auto const_op =
        dyn_cast_or_null<mlir::stablehlo::ConstantOp>(operand.getDefiningOp());
    if (!const_op) return failure();

    // Only support float tensors.
    auto tensor_type = mlir::dyn_cast_or_null<TensorType>(const_op.getType());