#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);

#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)

      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>(

    ],
    licenses = ["notice"],
)

package_group(
    name = "friends",
    packages = [
        "//tensorflow/compiler/mlir/tfr/...",
    ],
)

gen_op_libraries(
    name = "mnist_ops",
    src = "ops_defs.py",
    deps = [
        "//tensorflow:tensorflow_py",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:math_ops_gen",

      "tf_saved_model.session_initializer"() {initializers = [@init_func]} : () -> ()
      func.func @init_func() attributes {tf_saved_model.exported_names = ["init_func"], tf_saved_model.initializer_type = "init_op"} {
        func.return
      }
    }
  )mlir";

  Block block;
  ModuleOp module_op = ParseModuleOp(kModuleOpStr, block, ctx_);

  func::FuncOp init_func_op = GetInitializerFunction(

  // Test case: Variable removed for the initializer function whose
  // initializer_type is "init_op".
  "tf_saved_model.session_initializer"() { initializers = [@init] } : () -> ()

  func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"], tf_saved_model.initializer_type = "init_op"} {