}
  } else if (llvm::isa<TF::BatchMatMulV2Op>(op)) {
    if (const auto adj_y = op->getAttrOfType<BoolAttr>("adj_y")) {
      is_supported_affine_op = !adj_y.getValue();
    }
  } else if (llvm::isa<TF::MatMulOp>(op)) {
    if (const auto adj_y = op->getAttrOfType<BoolAttr>("transpose_b")) {
      is_supported_affine_op = !adj_y.getValue();
    }
  }

  // to stateful, otherwise the op will not be executed.
  OpBuilder builder(&ctx);
  module_op.walk([&builder](Operation* op) {
    if (op->hasAttrOfType<BoolAttr>("is_stateless") &&
        ContainCalibrationStatisticsSaverOp(op)) {
      op->setAttr("is_stateless", builder.getBoolAttr(false));
    }
  });
}

std::unique_ptr<OperationPass<ModuleOp>>

      });
      return false;
    }
    // FusedMatMul kernel does not support grad_a/grad_b attrs
    if ((matmul->hasAttr("grad_a") &&
         mlir::cast<BoolAttr>(matmul->getAttr("grad_a")).getValue()) ||
        (matmul->hasAttr("grad_b") &&
         mlir::cast<BoolAttr>(matmul->getAttr("grad_b")).getValue())) {
      (void)rewriter.notifyMatchFailure(matmul, [&](Diagnostic &diag) {

    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, res_type, min_value, max_value, quant_dim, num_bits,

    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, input_type, min_value, max_value, quant_dim, num_bits,

    This op returns a scalar string tensor as a key for user to look for the loaded array
    and a future containing the restored tensor.
  }];

  let arguments = (ins
    TF_Tensor:$variable,
    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
  );

  let results = (outs
    TF_Tensor:$array_key,
    MlrtFutureType: $tensor_future
  );
}

    } else if (auto int_attr =
                   mlir::dyn_cast_or_null<IntegerAttr>(attr.second)) {
      fbb.Int(attr.first.data(), int_attr.getInt());
    } else if (auto bool_attr = mlir::dyn_cast_or_null<BoolAttr>(attr.second)) {
      fbb.Bool(attr.first.data(), bool_attr.getValue());
    } else if (auto string_attr =
                   mlir::dyn_cast_or_null<StringAttr>(attr.second)) {

    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, res_type, min_value, max_value, quant_dim, num_bits,

      failed(AddFloatAttr(func, attrs, "h_scale", &fbb)) ||
      failed(AddFloatAttr(func, attrs, "w_scale", &fbb)))
    return failure();
  auto use_regular_nms =
      mlir::dyn_cast_or_null<BoolAttr>(attrs.get("use_regular_nms"));
  if (!use_regular_nms) {
    return func.emitError()
           << "use_regular_nms attribute is not set or not a bool";
  }
  fbb.Int("use_regular_nms", use_regular_nms.getValue());

    Other ops executed by TFRT may make use of $tensor_future.
  }];

  let arguments = (ins
    Arg<TF_Tensor, [{The variable tensor to be loaded}], []>:$variable,
    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
  );

  let results = (outs
    TF_StrTensor:$array_key,
    TF_Tensor: $tensor_future
  );

  TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>;