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

# CHECK: "tf_saved_model.global_tensor"() <{is_mutable, sym_name = "[[VAR0:[a-zA-Z_0-9]+]]", type = tensor<5x3xf32>, value = {{.*}} : tensor<5x3xf32>}> : () -> ()
# CHECK: "tf_saved_model.global_tensor"() <{is_mutable, sym_name = "[[VAR1:[a-zA-Z_0-9]+]]", type = tensor<3x5xf32>, value = {{.*}} : tensor<3x5xf32>}> : () -> ()
# CHECK:      func {{@[a-zA-Z_0-9]+}}(

    // CHECK: %[[VAR0:.*]] = "tf.VarHandleOp"()
    %0 = "tf.VarHandleOp"() {container = "", shared_name = "var0"} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>

    // CHECK: "tf.AssignVariableOp"(%[[VAR0]], %[[ASSET0]])
    "tf.AssignVariableOp"(%0, %arg0) : (tensor<!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()

    // CHECK: %[[VAR1:.*]] = "tf.VarHandleOp"()

    loop body inine):

    ```mlir
      %var0 = ...
      %var1 = ...
      tf.while (..., %var0, %var1) {
        tf_device.replicate ([%var0, %var1] as %rvar) {
          %compile:2 = "tf._TPUCompileMlir"()
          tf.TPUExecuteAndUpdateVariablesOp(%rvar, compile#1)
        }
      }
    ```

    This pass will transform it into

    ```mlir
      %var0 = ...
      %var1 = ...
      %state_var0 = ...

func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
  %var0 = "tf.VarHandleOp"() {container = "", shared_name = "var0"} : () -> tensor<!tf_type.resource<tensor<i1>>>
  %cond = "tf.ReadVariableOp"(%var0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i1>>>) -> tensor<i1>

         "      input_arg {"
         "        name: 'var0'"
         "        type: DT_RESOURCE"
         "      }"
         "      output_arg {"
         "        name: 'var0_value'"
         "        type: DT_FLOAT"
         "      }"
         "    }"
         "    node_def {"
         "      name: 'read0'"
         "      op: 'ReadVariableOp'"
         "      input: 'var0'"
         "      attr {"

loop body inine):

```mlir
  %var0 = ...
  %var1 = ...
  tf.while (..., %var0, %var1) {
    tf_device.replicate ([%var0, %var1] as %rvar) {
      %compile:2 = "tf._TPUCompileMlir"()
      tf.TPUExecuteAndUpdateVariablesOp(%rvar, compile#1)
    }
  }
```

This pass will transform it into

```mlir
  %var0 = ...
  %var1 = ...
  %state_var0 = ...
  %state_var1 = ...

    func.return %4 : tensor<100x50xf32>
  }
  // CHECK: "tf.VarHandleOp"
  // CHECK-SAME: shared_name = "var1"
  // CHECK-SAME: _is_initialized = true
  // CHECK: "tf.VarHandleOp"
  // CHECK-SAME: shared_name = "var2"
  // CHECK-SAME: _is_initialized = true
  // CHECK: "tf.VarHandleOp"

// CHECK-NEXT:  "tf.AssignVariableOp"([[VAR]], [[RES_1]])
// CHECK:       "tf.Yield"([[RES_3]], [[RES_4]])
func.func @readvariableop_is_not_hoisted_if_not_readonly(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  %cst_0 = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>

  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"
  // CHECK: %[[VAR:.*]] = "tf.VarHandleOp"()
  // CHECK-SAME: "var1"
  // CHECK: %[[CST:.*]] = arith.constant dense<> : tensor<0xf32>
  // CHECK: "tf.AssignVariableOp"(%[[VAR]], %[[CST]])
}


// -----

// Test invalid tensor
// expected-error@below{{Can't find variable var5 in session}}
module attributes {tf_saved_model.semantics, tf_saved_model.under_construction} {

        compile {
            attributes.attribute(buildType, 'debug')
            outgoing {
                variants {
                    var1 {
                        artifact oneJar
                        attributes.attribute(flavor, 'one')
                    }
                    var2 {
                        artifact twoJar
                        attributes.attribute(flavor, 'two')
                    }
                }