Searched refs:final_sort (Results 1 – 3 of 3) sorted by relevance
202 xla::XlaOp final_sort = xla::Sort( in DataOutputFastPath() local207 xla::XlaOp output = xla::GetTupleElement(final_sort, 1); in DataOutputFastPath()
1729 HloInstruction* final_sort = b_.AddInstruction(HloInstruction::CreateSort( in HandleSort() local1733 final_sort->set_sharding(HloSharding::Replicate() in HandleSort()1734 .GetTupleSharding(final_sort->shape()) in HandleSort()1736 PartitionedHlo replicated_sort(final_sort, final_sort->shape(), in HandleSort()
2455 auto final_sort = FindInstruction(module.get(), "sort.1"); in TEST_F() local2456 EXPECT_EQ(final_sort->operand(0)->shape().dimensions(1), 4000); in TEST_F()2457 EXPECT_EQ(final_sort->operand(1)->shape().dimensions(1), 4000); in TEST_F()2534 auto final_sort = FindInstruction(module.get(), "sort.1"); in TEST_F() local2535 EXPECT_EQ(final_sort->operand(0)->shape().dimensions(1), 4000); in TEST_F()2536 EXPECT_EQ(final_sort->operand(1)->shape().dimensions(1), 4000); in TEST_F()