class TestModule(tf.Module):

  def __init__(self):
    super(TestModule, self).__init__()
    self.child1 = Child()
    self.child2 = self.child1

  # CHECK: tf_saved_model.global_tensor
  # CHECK-SAME: tf_saved_model.exported_names = ["child1.my_variable", "child2.my_variable"]


if __name__ == '__main__':

  def __init__(self):
    super(TestModule, self).__init__()
    self.model = mnist_model()

  # CHECK: func {{@[a-zA-Z_0-9]+}}(%arg0: tensor<1x28x28x1xf32> {tf._user_specified_name = "x", tf_saved_model.index_path = [0]}
  # CHECK: attributes {{.*}} tf_saved_model.exported_names = ["my_predict"]
  @tf.function(input_signature=[
      tf.TensorSpec([1, 28, 28, 1], tf.float32),
  ])
  def my_predict(self, x):
    return self.model(x)

  }

  void runOnOperation() override {
    // Pass an invalid session argument, which is a nullptr.
    if (failed(mlir::tf_saved_model::MarkInitializedVariablesInFunction(
            getOperation(), /*session=*/nullptr)))
      return signalPassFailure();
  }
};

}  // namespace

namespace tf_saved_model {

static PassRegistration<MarkInitializedVariablesTestPass>
    mark_initialized_variables_test_pass;

    if (auto session_initializer =
            tf_saved_model::GetSessionInitializerOp(module)) {
      // clang-format off
      // Skip session initializer functions which are present in saved models.
      // For example:
      //   "tf_saved_model.session_initializer"() {initializers = [@init_all_tables]} : () -> () // NOLINT

  func.func @main(%arg0: tensor<10x1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<10x1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {

#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_INSERTCALLONCEOPFROMSESSIONINITIALIZERPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// This pass inserts a TFL::CallOnce op when tf_saved_model's session
// initializer is given.
class InsertCallOnceOpFromSessionInitializerPass

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/utils/name_utils.h"

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::GetSessionInitializerOp;
using ::mlir::tf_saved_model::kTfSavedModelExportedNamesAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

#include "mlir/IR/UseDefLists.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_mgr.h"

#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace tf_saved_model {
namespace {

using mlir::TF::VarHandleOp;

#define GEN_PASS_DEF_REMOVEVARIABLESINSESSIONINITIALIZERPASS

#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace tf_saved_model {

namespace {

#define GEN_PASS_DEF_ADDFUNCTIONSFOREXPORTEDNAMESPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_savedmodel_passes.h.inc"