// CHECK: [[RNG:%.*]] = "mhlo.rng"([[LOWER]], [[UPPER]], [[SHAPE]]) <{rng_distribution = #mhlo.rng_distribution<UNIFORM>}>
  // CHECK: [[SORT:%.*]]:2 = "mhlo.sort"([[RNG]], [[INPUT]]) <{dimension = -1 : i64, is_stable = {{.*}}}> ({
  // CHECK: ^{{.*}}([[ARG1:%.*]]: tensor<i32>, [[ARG2:%.*]]: tensor<i32>, {{.*}}: tensor<f32>, {{.*}}: tensor<f32>):
  // CHECK:   mhlo.compare LT, [[ARG1]], [[ARG2]], TOTALORDER

    "mhlo.return"(%3) : (tensor<i1>) -> ()
  }) {dimension = 1 : i64, is_stable = true} : (tensor<3x6xf32>, tensor<3x6xi32>) -> (tensor<3x6xf32>, tensor<3x6xi32>)
  func.return %2#0, %2#1 : tensor<3x6xf32>, tensor<3x6xi32>
}

// CHECK-LABEL:   func @convert_sort_to_topk_iotacst_broadcast(

        auto sorted = createSortOp(
            &rewriter, op.getLoc(), {keys, current},
            {rewriter.getIntegerType(32), input_type.getElementType()},
            /*dimension=*/-1, /*isStable=*/false,
            /*direction=*/ComparisonDirection::LT);
        current = sorted.getResult(1);
      }
      rewriter.replaceOp(op, current);
      return success();
    }