body_subgraph_index : int;
}

table StablehloWhileOptions{
  cond_subgraph_index : int;
  body_subgraph_index : int;
}

table StablehloSortOptions{
  dimension : long;
  is_stable : bool;
  comparator_subgraph_index : int;
}

table StablehloConcatenateOptions {
  dimension : long;
}

table StablehloBroadcastInDimOptions{
  broadcast_dimensions : [long];
}

        return when {
            this is FirSmartCastExpression -> {
                val result = originalExpression.toKtReceiverValue()
                if (result != null && isStable) {
                    KaSmartCastedReceiverValue(result, smartcastType.coneType.asKtType())
                } else {
                    result
                }
            }

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

  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