#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
#include "tensorflow/python/lib/core/pybind11_status.h"

PYBIND11_MODULE(mlir_wrapper, m) {
  m.def("preloadTensorFlowDialects", [](mlir::MLIRContext &context) {
    mlir::DialectRegistry registry;
    mlir::RegisterAllTensorFlowDialects(registry);
    context.appendDialectRegistry(registry);
    context.loadAllAvailableDialects();
  });

    def get(self) -> UnrankedTensorType: ...

class Value:
    def __init__(self, *args, **kwargs) -> None: ...
    def getType(self, *args, **kwargs) -> Any: ...

def preloadTensorFlowDialects(arg0) -> None: ...