"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(

  void add_dimension(int64_t dimension) {
    fbb_.AddElement<int64_t>(StablehloSortOptions::VT_DIMENSION, dimension, 0);
  }
  void add_is_stable(bool is_stable) {
    fbb_.AddElement<uint8_t>(StablehloSortOptions::VT_IS_STABLE, static_cast<uint8_t>(is_stable), 0);
  }
  void add_comparator_subgraph_index(int32_t comparator_subgraph_index) {

        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();
    }

is_stable attributes.
  }];

  let arguments = (ins
    Arg<Variadic<TF_Tensor>, [{A list of `Tensor` of identical shape but possibly different types.}]>:$inputs,
    Arg<TF_Int32Tensor, [{The dimension along which to sort. Must be a compile-time constant.}]>:$dimension,

    SymbolRefAttr:$comparator,
    BoolAttr:$is_stable
  );

  let results = (outs