"MatrixDiag",
    "MatrixDiagPart",
    "MatrixDiagPartV2",
    "MatrixDiagPartV3",
    "MatrixDiagV2",
    "MatrixDiagV3",
    "MatrixSetDiag",
    "MatrixSetDiagV2",
    "MatrixSetDiagV3",
    "MirrorPad",
    "OneHot",
    "OnesLike",
    "Pack",
    "Pad",
    "PadV2",
    "ParallelConcat",
    "PlaceholderV2",
    "PlaceholderWithDefault",
    "PreventGradient",

            (Arith_ConstantOp ConstantAttr<I32ElementsAttr, "{0}">)),
          (TF_MatrixSetDiagOp $input, $diagonal)>;

// `align` attribute can be ignored because we only support converting
// `MatrixSetDiagV3` to `MatrixSetDiag` with default `k` inputs.
def ConvertMatrixSetDiagV3ToMatrixSetDiag :
      Pat<(TF_MatrixSetDiagV3Op $input, $diagonal,
            (Arith_ConstantOp ConstantAttr<I32ElementsAttr, "{0}">),

  %cst = arith.constant dense<1> : tensor<i32>
  %0 = "tf.MatrixSetDiagV3"(%arg0, %arg1, %cst) : (tensor<3x3xi64>, tensor<3xi32>, tensor<i32>) -> tensor<3x3xi64>
  func.return %0 : tensor<3x3xi64>

// CHECK-LABEL: @matrix_set_diag_v3_non_zero_k
// CHECK: tf.MatrixSetDiagV3
}

func.func @matrix_set_diag_v3_default_align(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi32>) -> tensor<3x3xi64> {

  func.return %0 : tensor<3x3xi64>

  // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}>
  // CHECK: %[[RES:.*]] = "tf.MatrixSetDiagV3"(%arg0, %arg1, %[[ZERO]])
  // CHECK-SAME: <{align = "RIGHT_LEFT"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"}
  // CHECK-SAME: (tensor<3x3xi64>, tensor<3xi64>, tensor<i32>) -> tensor<3x3xi64>
}

      "MatrixDiagPartV2",
      "MatrixDiagPartV3",
      "MatrixDiagV2",
      "MatrixDiagV3",
      "MatrixInverse",
      "MatrixSetDiag",
      "MatrixSetDiagV2",
      "MatrixSetDiagV3",
      "MatrixSolve",
      "MatrixTriangularSolve",
      "MaxPool",
      "MaxPool3D",
      "MaxPool3DGrad",
      "MaxPool3DGradGrad",
      "MaxPoolGrad",
      "MaxPoolGradGrad",

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;

  let hasCanonicalizer = 1;
}

def TF_MatrixSetDiagV3Op : TF_Op<"MatrixSetDiagV3", [Pure]> {
  let summary = [{
Returns a batched matrix tensor with new batched diagonal values.
  }];

  let description = [{
Given `input` and `diagonal`, this operation returns a tensor with the