dyn_cast<arith::ConstantOp>(op.getOperand(0).getDefiningOp());
    if (!input_op) {
      return failure();
    }

    Builder builder(op.getContext());
    auto new_gather_op = rewriter.create<TFL::GatherOp>(
        op.getLoc(),
        /*result=*/
        mlir::cast<TensorType>(op.getResult().getType())
            .clone(builder.getI4Type()),
        /*operand=*/op.getOperands(), op->getAttrs());

                            assign.getValue().getType());
  }

  OpBuilder builder(replicate_op);
  // Clone this old replicate op but with new result types.
  auto new_replicate_op = builder.create<tf_device::ReplicateOp>(
      replicate_op->getLoc(), new_result_types, replicate_op->getOperands(),
      replicate_op->getAttrs());

  // Move region to the new op.

    const SmallVector<mlir::Type, NOutput>& output_types) {
  auto func_type = builder->getFunctionType(input_types, output_types);
  auto func = func::FuncOp::create(
      mlir::NameLoc::get(builder->getStringAttr("fused_func")), "fused_func",
      func_type, {});

  func.addEntryBlock();
  mlir::StringAttr attr_value = builder->getStringAttr("MaxUnpooling2D");

    } else {
      /* If we're not running on a TPU node, we might not be able to
       * actually call the part of the TPU API that gives us layout.
       * This happens e.g. for unit tests. Below we just create a reasonable
       * layout.  We sort by dimension size, which makes the layout agree with
       * the "correct" TPU layout in surprisingly many cases.
       * Note that the corresponding InfeedEnqueue op will be generated

using tensorflow::tracing::TracingTensorHandle;

void TF_SetTracingImplementation(const char* name, TF_Status* s) {
  tsl::Set_TF_Status_from_Status(s, SetDefaultTracingEngine(name));
}

// Creates a new TensorFlow function, it is an execution context attached to a
// given tracing context.
TF_ExecutionContext* TF_CreateFunction(const char* fn_name, TF_Status* s) {
  return wrap(CreateTracingExecutionContext(fn_name, s));

  // reachable from the `fetch` operands can be eliminated.

  llvm::SmallPtrSet<Operation*, 8> reachable_ops;
  llvm::SmallVector<Operation*, 8> ops_to_visit;

  // Visit fetches first to create a starting point for ops that are reachable.
  reachable_ops.insert(graph.GetFetch());
  VisitOpOperands(graph, graph.GetFetch(), &reachable_ops, &ops_to_visit);

    auto attr =
        builder.getIntegerAttr(builder.getIntegerType(64), arithmetic_count);
    op->setAttr("_arithmetic_count", attr);
  });
}

}  // namespace

/// Creates an instance of the TensorFlow Lite dialect GetArithmeticCount
/// pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateGetArithmeticCountPass() {
  return std::make_unique<GetArithmeticCountPass>();
}

    // of `scf.if` from that of `tf.IfRegion`.
    auto scf_if_condition =
        rewriter.create<tensor::ExtractOp>(loc, op.getCond());

    TypeRange tf_if_region_return_type = op.getResultTypes();

    // Create the `scf.if` op.
    auto scf_if_op =
        rewriter.create<scf::IfOp>(loc, tf_if_region_return_type,
                                   scf_if_condition, /*withElseRegion=*/true);

    InsertTFLQuantOpsAfterTFFakeQuantOp<
        TF::FakeQuantWithMinMaxArgsOp, /*PerAxis=*/false,
        FetchMinMaxAttrs<TF::FakeQuantWithMinMaxArgsOp>>;

// Removes the wrapper of the tf.FakeQuant* ops and creates the tfl.quantize
// and tfl.dequantize pairs before tf.FakeQuant* being foled.
LogicalResult ConvertFakeQuantOps(func::FuncOp func, MLIRContext* ctx,
                                  bool use_fake_quant_num_bits) {

    custom_op->setAttrs(inner_op->getAttrs());
    builder.create<YieldOp>(loc, inner_op->getResults());
    custom_op.getBody().takeBody(region);

    op->replaceAllUsesWith(custom_op);
    op->erase();
  }
}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass() {
  return std::make_unique<RaiseCustomOpsPass>();
}

// Creates an instance of the TensorFlow Lite dialect raise custom op pass.