"$_builder.getStringAttr(\"" # x # "\")">;

def MatrixSetDiagToV3 : Pat<
  (TF_MatrixSetDiagOp:$src $input, $diag),
  (TF_MatrixSetDiagV3Op:$dest $input, $diag, (TF_ConstOp (GetI32Attr<0>)),
    (GetStrAttr<"RIGHT_LEFT">)),
  [], [(CopyAttrs $src, $dest)]>;

// MatrixSetDiagToV2 op implicitly used LEFT_LEFT alignment.
def MatrixSetDiagV2ToV3 : Pat<
  (TF_MatrixSetDiagV2Op:$src $input, $diag, $k),

    // CHECK: mhlo.reshape
    // CHECK: mhlo.concatenate
    // CHECK: mhlo.gather
    // CHECK: mhlo.broadcast
    // CHECK: mhlo.select
    %2 = "tf.MatrixDiagPartV3"(%arg0, %1, %0) {
        T = i32, align = "RIGHT_LEFT"
    } : (tensor<7x140x128xi32>, tensor<2xi32>, tensor<i32>) -> tensor<7x22x128xi32>
    func.return %2: tensor<7x22x128xi32>
  }

  // CHECK-LABEL: func @xla_call_module

  // CHECK-DAG: %[[ZEROF:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
  // CHECK-DAG: "tf.MatrixDiagV3"(%arg0, %[[ZEROI]], %[[MINUS1]], %[[MINUS1]], %[[ZEROF]]) <{align = "RIGHT_LEFT"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x4xf32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<f32>) -> tensor<2x4x4xf32>

  %0 = mhlo.constant dense<42> : tensor<i32>  // padding value
  %1 = mhlo.constant dense<[-10, 11]> : tensor<2xi32>  // k
  %2 = "tf.MatrixDiagPartV3"(%arg0, %1, %0) {
      T = i32, align = "RIGHT_LEFT"
  } : (tensor<7x140x128xi32>, tensor<2xi32>, tensor<i32>) -> tensor<7x22x128xi32>
  func.return %2: tensor<7x22x128xi32>
}

// -----

// CHECK-LABEL: func @matrix_diag_part_zero_dim_complex

    // left, "RIGHT" means rows will be padded ot the right.  The default is
    // "RIGHT_LEFT".
    StringRef align = op->getAttrOfType<StringAttr>("align").getValue();
    enum Alignment { kLeft, kRight };

    // default is RIGHT_LEFT
    Alignment superdiagonal_align = kRight;
    Alignment subdiagonal_align = kLeft;

    if (align == "RIGHT_LEFT") {
      superdiagonal_align = kRight;
      subdiagonal_align = kLeft;

    Arg<TF_Tensor, [{The value to fill the area outside the specified diagonal band with.
Default is 0.}]>:$padding_value,

    DefaultValuedOptionalAttr<TF_AnyStrAttrOf<["LEFT_RIGHT", "RIGHT_LEFT", "LEFT_LEFT", "RIGHT_RIGHT"]>, "\"RIGHT_LEFT\"">:$align
  );

  let results = (outs
    Res<TF_Tensor, [{The extracted diagonal(s).}]>:$diagonal
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

}

func.func @matrix_set_diag_v3(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi32>) -> tensor<3x3xi64> {
  %cst = arith.constant dense<0> : tensor<i32>
  %0 = "tf.MatrixSetDiagV3"(%arg0, %arg1, %cst) {align = "RIGHT_LEFT"} : (tensor<3x3xi64>, tensor<3xi32>, tensor<i32>) -> tensor<3x3xi64>
  func.return %0 : tensor<3x3xi64>

// CHECK-LABEL: func @matrix_set_diag_v3