namespace cpp {

ArgView::ArgView(ArgSpec arg) : arg_(arg) {}

string ArgView::VariableName() const { return arg_.name(); }

string ArgView::SetterMethod() const {
  if (IsList()) {
    return "AddInputList";
  } else {
    return "AddInput";
  }
}

std::vector<string> ArgView::SetterArgs() const { return {VariableName()}; }

bool ArgView::IsList() const { return arg_.arg_type().is_list(); }

namespace tensorflow {
namespace generator {
namespace cpp {

class ArgView {
 public:
  explicit ArgView(ArgSpec arg);

  string VariableName() const;
  string SetterMethod() const;
  std::vector<string> SetterArgs() const;
  int Position() const;

  bool IsList() const;

 private:
  ArgSpec arg_;
};

}  // namespace cpp
}  // namespace generator
}  // namespace tensorflow

//
// - arg_Rt_31_1__W_0__X_1:
//     a W or X register encoded in Rt[4:0]. The width specifier is encoded in the field
//     [31:31] (offset 31, bit count 1) and the register is W for 0 and X for 1.
//
// - arg_[s|u]label_FIELDS_POWER:
//     a program label encoded as "FIELDS" times 2^POWER in the range [MIN, MAX] (determined
//     by signd/unsigned, FIELDS and POWER), e.g.
//       arg_slabel_imm14_2
//       arg_slabel_imm19_2

func.func @testUnidirectionalSequenceLstmShapeInference(%arg0: tensor<600 x 10 x 20 x f32>, %arg1: tensor<? x ? x f32>, %arg2: tensor<? x ? x f32>, %arg3: tensor<? x ? x f32>, %arg4: tensor<? x ? x f32>, %arg5: tensor<? x ? x f32>, %arg6: tensor<? x ? x f32>, %arg7: tensor<? x ? x f32>, %arg8: tensor<? x ? x f32>, %arg9: tensor<? x f32>, %arg10: tensor<? x f32>, %arg11: tensor<? x f32>, %arg12: tensor<? x f32>, %arg13: tensor<? x f32>,...

  // Create xla_params.
  std::vector<xla::XlaOp> xla_params;
  for (mlir::BlockArgument& arg : block.getArguments()) {
    auto num = arg.getArgNumber();
    xla::Shape shape = xla::TypeToShape(arg.getType());
    xla::XlaOp argop =
        xla::Parameter(&builder, num, shape, absl::StrCat("Arg_", num));
    xla_params.push_back(argop);
  }

  std::vector<xla::XlaOp> returns(1);

  // Test using the argN() methods.
  {
    matmul.arg0(0, 0) = 1;
    matmul.arg0(0, 1) = 2;
    matmul.arg0(0, 2) = 3;
    matmul.arg0(1, 0) = 4;
    matmul.arg0(1, 1) = 5;
    matmul.arg0(1, 2) = 6;

    matmul.arg1(0, 0) = 7;
    matmul.arg1(0, 1) = 8;
    matmul.arg1(1, 0) = 9;
    matmul.arg1(1, 1) = 10;
    matmul.arg1(2, 0) = 11;
    matmul.arg1(2, 1) = 12;

  // CHECK: "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]], %[[ARG_4]], %[[ARG_5]], %[[CONST_0]], %[[CONST_0]], %[[CONST_0]], %[[CONST_MODE]])

// Ensure lstm roundtrip exactly

// CHECK-SAME:  ([[ARG0:%.+]]: tensor<i1>, [[ARG1:%.+]]: tensor<i32>, [[ARG2:%.+]]: tensor<i64>, [[ARG3:%.+]]: tensor<f32>, [[ARG4:%.+]]: tensor<f64>, [[ARG5:%.+]]: tensor<!tf_type.resource>, [[ARG6:%.+]]: tensor<!tf_type.resource>)
func.func @replace_multiple_outputs(%arg0: tensor<i1>, %arg1: tensor<i32>, %arg2: tensor<i64>, %arg3: tensor<f32>, %arg4: tensor<f64>, %arg5: tensor<!tf_type.resource>, %arg6: tensor<!tf_type.resource>) {

// CHECK-LABEL: func.func private @"map/while_body/TfMlrtAsyncWhileBody"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor<i32>, %arg6: tensor<?x!tf_type.resource>, %arg7: tensor<*xf32>) {
// CHECK-NEXT:    %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT:    %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>