| 455 | MGB_DYN_TYPE_OBJ_FINAL_IMPL(TopK); |
| 456 | |
| 457 | TopK::TopK( |
| 458 | VarNode* data, VarNode* k, const Param& param, const OperatorNodeConfig& config) |
| 459 | : Super(data->owner_graph(), config, "top_k", {data, k}) { |
| 460 | init_megdnn_opr(*this, param); |
| 461 | add_input({data, k}); |
| 462 | if (param.mode == Param::Mode::KTH_ONLY) { |
| 463 | output(1) |
| 464 | ->add_flag(VarNode::Flag::VOLATILE_CONTENT) |
| 465 | .add_flag(VarNode::Flag::ALLOW_EMPTY_SHAPE); |
| 466 | } |
| 467 | } |
| 468 | |
| 469 | std::array<SymbolVar, 2> TopK::make( |
| 470 | SymbolVar data, SymbolVar k, const Param& param, |