#include "pybind11/stl.h"  // from @pybind11

namespace py = pybind11;

void init_basic_classes(py::module& m);
void init_types(py::module& m);
void init_builders(py::module& m);
void init_ops(py::module& m);
void init_attrs(py::module& m);

      module->emitError("Invalid MLIR module: failed verification.");
      return false;
    }
    return true;
  });

  init_basic_classes(m);
  init_types(m);
  init_builders(m);
  init_ops(m);
  init_attrs(m);

// Get a set of node names that must be preserved. They can not be transformed
// or removed during the graph transformation. This includes feed and fetch
// nodes, keep_ops, init_ops. Fills in `num_values` and `storage_size`, they
// will be used in `TF_GetNodesToPreserveList`.
TF_CAPI_EXPORT extern void TF_GetNodesToPreserveListSize(
    const TF_GrapplerItem* item, int* num_values, size_t* storage_size,

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

void init_ops(py::module& m) {
  py::class_<mlir::Operation, std::unique_ptr<mlir::Operation, py::nodelete>>(
      m, "Operation")
      .def("getRegion", &mlir::Operation::getRegion,
           py::return_value_policy::reference)

TEST_F(ConvertMlirModuleToExportedModelTest, InitNodeNotSetIfLocNameMismatch) {
  // Define a module that initializes a tf.HashTableV2 whose control output node
  // for the initialization is named "init_ok". Since the output control node
  // name does not begin with "init_op" the init node could not have been found
  // after the conversion.
  mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(

// CHECK-LOC: tf_executor.island({{.*}}) wraps "tf.NoOp"()
// CHECK-LOC-SAME: loc("init_op_NoOp")
}

// -----

// Test the case where there are 2 initializer functions ("init_op" and
// "restore_op"). The init func of type "init_op" is merged first.

// CHECK-LABEL: module attributes
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1228 : i32}, tf_saved_model.semantics} {

      for (auto init_op : init_table_from_text_file_ops_to_erase) {
        // Replace the InitializeTableFromTextFileV2Op to use the saved model's
        // asset filepath.
        builder.setInsertionPoint(init_op);
        builder.create<TF::InitializeTableFromTextFileV2Op>(
            init_op.getLoc(), init_op.getTableHandle(), const_op.getResult(),
            init_op.getKeyIndex(), init_op.getValueIndex(),

    signature_def_map, init_op, assets_collection = create_signature()

    sess = tf.Session()
    sess.run(tf.initializers.global_variables())
    builder = tf.saved_model.builder.SavedModelBuilder(save_model_path)
    builder.add_meta_graph_and_variables(
        sess,
        [tf.saved_model.tag_constants.SERVING],
        signature_def_map,
        main_op=init_op,
        assets_collection=assets_collection,

  y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
  math_ops.add(x, y, name='x_y_sum')

  init_op = variables.global_variables_initializer()
  saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V1)
  with session.Session() as sess:
    sess.run(init_op)
    sess.run(y.assign(y + 42))
    # Without the checkpoint, the variable won't be set to 42.

    ModuleOp module = getOperation();
    if (!HasTfSavedModelSemantics(module)) {
      return;
    }

    auto init_op = GetSessionInitializerOp(module);
    if (init_op == nullptr || init_op.getInitializers().empty()) {
      return;
    }

    SymbolTable symbol_table(module);
    for (auto initializer : init_op.getInitializers()) {
      auto func = symbol_table.lookup<func::FuncOp>(