# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-graph-as-function -o - | FileCheck %s

node {
  name: "arg"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {
  name: "ret"
  op: "_Retval"
  input: "arg"
  attr {
    key: "T"
    value {
      type: DT_INT32

      }
    }
  }
}
node {
  name: "input_1_1"
  op: "_Arg"
  device: "/job:localhost/replica:0/task:0/device:GPU:0"
  attr {
    key: "T"
    value {
      type: DT_HALF
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {
  name: "bn2a_branch2c_cond_input_1"
  op: "_Arg"
  device: "/job:localhost/replica:0/task:0/device:GPU:0"
  attr {

    tf_executor.fetch %4#0 : tensor<10xi32>
  }
  func.return %graph : tensor<10xi32>
}

// CHECK: name: "foo"
// CHECK-NEXT: op: "_Arg"
// CHECK: name: "bar"
// CHECK-NEXT: op: "_Arg"
// CHECK: name: "[[BAR_ID_0:.*]]"
// CHECK-NEXT: op: "Identity"
// CHECK-NEXT: input: "bar"
// CHECK: name: "[[BAR_ID_1:.*]]"
// CHECK-NEXT: op: "Identity"
// CHECK-NEXT: input: "[[BAR_ID_0]]"

node {
  name: "args_0"
  op: "_Arg"
  device: "/CPU:0"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {
  name: "args_1"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {

  auto arg0 = ops::_Arg(scope.WithOpName("a_0_arg"), DT_INT32, 0);
  auto arg1 = ops::_Arg(scope.WithOpName("b_0_arg"), DT_FLOAT, 1);
  auto arg2 = ops::_Arg(scope.WithOpName("c_0_arg"), DT_INT32, 2);
  auto arg3 = ops::_Arg(scope.WithOpName("d_0_arg"), DT_FLOAT, 3);

  auto arg4 = ops::_Arg(scope.WithOpName("u_0_arg"), DT_RESOURCE, 4);
  auto arg5 = ops::_Arg(scope.WithOpName("v_0_arg"), DT_RESOURCE, 5);

  // TODO(b/74182462): We implement this by creating a new dummy Graph including
  // _Arg nodes, and let CompileGraph walk it. This could be optimized.
  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));

  // First create the actual node we care about computing.
  TF_ASSIGN_OR_RETURN(Node * main_node, graph->AddNode(node_def));

  // Create dummy _Arg nodes. Link these to `node` and also via a control

     tf_executor.fetch %outputs : tensor<8x1xf32>
  }
  func.return %0 : tensor<8x1xf32>
}

// Check that we generate _handle_dtypes and _handle_shapes for the resource
// argument.

// CHECK:      op: "_Arg"

// CHECK:        key: "_handle_dtypes"
// CHECK-NEXT:   value {
// CHECK-NEXT:     list {
// CHECK-NEXT:       type: DT_FLOAT
// CHECK-NEXT:     }
// CHECK-NEXT:   }
// CHECK-NEXT: }

}

TEST(CreateSaverDefTest, CreateValidSaverDef) {
  // Needs to have a _Arg node with an attribute "tf_saved_model.index_path" =
  // ["__tf_file_prefix"].
  GraphDef graph_def;
  ASSERT_TRUE(TextFormat::ParseFromString(
      R"pb(node {
             name: "foo",
             op: "_Arg",
             attr {
               key: "tf_saved_model.index_path",

Clustered nodes: 2725
Unclustered nodes: 606
Number of clusters: 2

unclustered size 606
 AssignAddVariableOp 1
 Const 108
 DivNoNan 1
 Identity 2
 Merge 53
 PlaceholderWithDefault 1
 ReadVariableOp 2
 Switch 1
 _Arg 435
 _Retval 2
cluster 0 size 1910
 Add 16
 AddN 71
 ArgMax 1
 AssignAddVariableOp 1
 BiasAdd 1
 BiasAddGrad 1
 Cast 115
 Const 407
 Conv2D 53
 Conv2DBackpropFilter 53
 Conv2DBackpropInput 52

node {
  name: "arg"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
versions {
  producer: 27
}