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


class TestModule(tf.Module):

  def __init__(self):
    super(TestModule, self).__init__()
    # CHECK: "tf_saved_model.global_tensor"() <{is_mutable, {{.*}} type = tensor<*xf32>, value = dense<0.000000e+00> : tensor<1xf32>}> {tf_saved_model.exported_names = ["v0"]} : () -> ()

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 975 : i32}, tf_saved_model.semantics}  {
  func.func @serving_default(%arg0: tensor<1x20x20x28xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x40x40x28xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {
      %c = stablehlo.constant dense<40> : tensor<2xi32>

// CHECK-NOT: tf_saved_model.semantics
// CHECK: tfrt.cost_threshold = 1024
// CHECK-SAME: tfrt.merge_inter_dependent_streams = true
module attributes {tf_saved_model.semantics} {

// CHECK-NOT: "tf_saved_model.session_initializer"
"tf_saved_model.session_initializer"() { initializers = [@func_init] } : () -> ()

// CHECK-LABEL: _tfrt_resource_init

  func.func @serving_default(%arg0: tensor<1024x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1024x3xf32> {tf_saved_model.index_path = ["output1"]}, tensor<1024x3xf32> {tf_saved_model.index_path = ["output2"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {

    (i.e. `obj[0]["x"]`).

    The function arg attribute
    `tf_saved_model.bound_input = @some_symbol_name` references a
    `tf_saved_model.global_tensor` that is bound to that function argument
    and implicitly passed in when the function is invoked from a user-level API.
    TODO(silvasean): Consider replacing tf_saved_model.bound_input arg attrs
    with "get_global @some_global_tensor" in the function body.
  }];

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

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

class DedupBoundInputBindingPass

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 975 : i32}, tf_saved_model.semantics}  {
  func.func @serving_default(%arg0: tensor<1x32x32x128xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x64x64x128xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {

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

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

    // Pass an invalid session argument, which is a nullptr.
    if (failed(tf_saved_model::LiftVariables(module, /*session=*/nullptr)))
      signalPassFailure();
  }
};

}  // namespace
}  // namespace tf_saved_model

namespace tf_test {

std::unique_ptr<OperationPass<ModuleOp>> CreateLiftVariablesTestPass() {
  return std::make_unique<tf_saved_model::LiftVariablesTestPass>();
}

std::unique_ptr<OperationPass<ModuleOp>>

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace mlir::quant {
namespace {

using ::mlir::tf_saved_model::AssetOp;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::LookupBoundInputOfType;
using ::tensorflow::AssetFileDef;