Searched refs:TopkRewriter (Results 1 – 4 of 4) sorted by relevance
26 class TopkRewriter : public HloModulePass {28 explicit TopkRewriter(std::function<bool(const HloSortInstruction*, int64)> in TopkRewriter() function
80 TopkRewriter rewriter([](const HloSortInstruction*, int64) { return true; }); in TEST_F()110 TopkRewriter rewriter([](const HloSortInstruction*, int64) { return true; }); in TEST_F()140 TopkRewriter rewriter([](const HloSortInstruction*, int64) { return true; }); in TEST_F()
66 StatusOr<bool> TopkRewriter::Run(HloModule* module) { in Run()
374 pipeline.AddPass<TopkRewriter>([](const HloSortInstruction* sort, int64) { in RunHloPassesThroughLayoutAssn()