tensorflow::Var** maybe_resource) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  *maybe_resource = nullptr;
  if (cc_ctx->input_dtype(input) == tensorflow::DT_RESOURCE) {
    if (LookupResource(cc_ctx, HandleFromInput(cc_ctx, input), maybe_resource)
            .ok()) {
      return (*maybe_resource)->mu();
    } else {

absl::StatusOr<ElementsAttr> ConvertTensor(const Tensor& input_tensor,
                                           Builder* builder) {
  const auto& input_dtype = input_tensor.dtype();
  const auto& input_shape = input_tensor.shape();
  Type elt_type;
  TF_RETURN_IF_ERROR(ConvertDataType(input_dtype, *builder, &elt_type));
  SmallVector<int64_t, 4> shape;
  ConvertToMlirShape(input_shape, &shape);

  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  CHECK_GE(index, 0);                     // Crash OK
  CHECK_LT(index, cc_ctx->num_inputs());  // Crash OK
  return static_cast<TF_DataType>(cc_ctx->input_dtype(index));
}

void TF_SetOutput(TF_OpKernelContext* ctx, int i, const TF_Tensor* tensor,
                  TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);

  auto input_type = RankedTensorType::get({1, 2, 2, 1}, builder_->getF32Type());
  auto output_type =
      RankedTensorType::get({1, 2, 1, 1}, builder_->getF32Type());
  SmallVector<mlir::Type, 1> input_types{input_type};
  SmallVector<mlir::Type, 1> output_types{output_type};

  auto max_unpooling_func =
      createMaxUnpoolingFunc<1, 1>(builder_.get(), input_types, output_types);

  }

  std::vector<Output> func_inputs;
  std::vector<DataType> input_dtypes;
  const int num_inputs = op.num_inputs();
  func_inputs.reserve(num_inputs + grad_inputs.size());
  input_dtypes.reserve(num_inputs);

  for (int i = 0; i < num_inputs; i++) {
    func_inputs.push_back(op.input(i));
    input_dtypes.push_back(op.input_type(i));
  }

  func_inputs.insert(std::end(func_inputs), std::begin(grad_inputs),

// process.
struct ExpressedToQuantizedConverter {
  // Creates a converter for the given input type.
  static ExpressedToQuantizedConverter forInputType(Type input_type);

  // Converts the inputType to be based on the given elemental type,
  // returning the new type (or nullptr and emit an error on failure).
  Type convert(quant::QuantizedType elemental_type) const;

  // Whether the conversion is legal.

  mlir::OpBuilder builder(func.getBody());
  auto& bb = func.front();
  auto loc = func.getLoc();

  int num_args = bb.getNumArguments();
  llvm::SmallVector<Type, 4> input_types;
  input_types.reserve(num_args);
  // Edit the block arguments and create the new input ops in place to replace
  // the old input ops and quantize ops.
  for (int i = 0; i != num_args; ++i) {

        !result_type.isa<IntegerType>()) {
      return nullptr;
    }
    auto input_itype = input.getType().cast<IntegerType>();
    auto result_itype = result_type.cast<IntegerType>();
    if (input_itype.getWidth() == result_itype.getWidth()) return nullptr;
    if (input_itype.getWidth() > result_itype.getWidth()) {
      return builder.create<arith::TruncIOp>(conversion_loc, result_type,

  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
  SmallVector<mlir::Type, 4> input_types{input_type, weight_type, bias_type,
                                         projection_type,
                                         layer_norm_scale_type};
  auto func_type = builder->getFunctionType(input_types, output_type);

  auto func = func::FuncOp::create(

                                      Operation** result) {
  auto input_type = mlir::cast<RankedTensorType>(input.getType());
  SmallVector<int64_t, 4> output_shape;
  int size_of_splits;
  if (input_type.getRank() < axis || axis < 0) return failure();
  for (int i = 0; i < input_type.getRank(); ++i) {
    int64_t dim = input_type.getDimSize(i);
    if (i == axis) {
      if (dim % splits != 0) {