}
};

// Converts TensorFlow DiagPartOp to HLO ops using reduction on masked matrix.
// For a Rank-2 input, it creates the following ops:
//   %1 = "mhlo.iota"() {iota_dimension = 0 : i64}
//   %2 = "mhlo.iota"() {iota_dimension = 1 : i64}
//   %3 = "mhlo.compare"(%1, %2) {comparison_direction = "EQ"}
//   %4 = mhlo.constant dense<0.000000e+00> : tensor<f32>
//   %5 = "mhlo.broadcast"(%4)
//   %6 = "mhlo.select"(%3, %input, %5)