#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

namespace mlir {
namespace tf_saved_model {

namespace {

std::string GetVariableName(Operation* op) {
  if (auto handle = dyn_cast<TF::VarHandleOp>(op)) {

// RUN: tf-tfrt-opt -tf-to-tfrt %s | FileCheck %s --dump-input=fail

// CHECK-LABEL: func @device_test
func.func @device_test(
    %arg0: tensor<3x1xf32> {tf_saved_model.index_path = [0]},
    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"

from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1

# Tests handling dupliate functions after V1 control flow is functionalized.

# CHECK:  func {{.*}} tf_saved_model.exported_names = ["key_1"]
# CHECK: "tf.If"
# CHECK-SAME: else_branch = @[[else:[a-zA-Z_0-9]+]]
# CHECK-SAME: then_branch = @[[then:[a-zA-Z_0-9]+]]

# CHECK:  func {{.*}} tf_saved_model.exported_names = ["key_2"]
# CHECK: "tf.If"

from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1

# Tests different SignatureDef's with identical method_name string

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME:   {{.*}})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME:   {{.*}})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key2"]


def Test():

// RUN: tf-tfrt-opt -pass-pipeline='builtin.module(tf-executor-to-tfrt-pipeline{target-tpurt=true})' %s | FileCheck %s

module attributes {tf_saved_model.semantics}  {
  // CHECK-LABEL: func @main
  func.func @main_func() -> (tensor<*xf32> {tf_saved_model.index_path = ["a"]}) attributes {tf_saved_model.exported_names = ["main_func"]} {
    %0 = tf_executor.graph {

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 554 : i32}, tf_saved_model.semantics} {

#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/ir/tf_types.h"
#include "tsl/platform/path.h"

namespace mlir {
namespace tf_saved_model {
namespace {

#define GEN_PASS_DEF_ASSETSINKINGPASS
#define GEN_PASS_DECL_ASSETSINKINGPASS

// CHECK-NEXT: }

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 554 : i32}, tf_saved_model.semantics} {
  func.func @main() -> (tensor<5xf32> {tf_saved_model.index_path = ["start_logits"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "StatefulPartitionedCall:1"}, tf_saved_model.exported_names = ["serving_default"]} {
    %cst = arith.constant dense<0.000000e+00> : tensor<5xf32>

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace tf_saved_model {

namespace {

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

import tensorflow.compat.v2 as tf
from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common


class ReferencesParent(tf.Module):

  def __init__(self, parent):
    super(ReferencesParent, self).__init__()
    self.parent = parent
    # CHECK: tf_saved_model.global_tensor
    # CHECK-SAME: tf_saved_model.exported_names = ["child.my_variable"]
    self.my_variable = tf.Variable(3.)