Function TEST_F

tensorflow/compiler/xla/service/stable_sort_expander_test.cc:83–112

}

TEST_F(StableSortExpanderTest, StabilizeSortReuseIotaOperand) {
  const char* hlo_string = R"(
   HloModule permutation_sort

   compare {
     p.0.lhs = f32[] parameter(0)
     p.0.rhs = f32[] parameter(1)
     p.1.lhs = s32[] parameter(2)
     p.1.rhs = s32[] parameter(3)
     ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT
   }

   ENTRY sort_computation {
     keys = f32[64,8732]{1,0} parameter(0)
     values = s32[64,8732]{1,0} iota(), iota_dimension=1
     sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values),
       dimensions={1}, to_apply=compare, is_stable=true
     ROOT gte = f32[64,8732]{1,0} get-tuple-element(sort), index=0
   })";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));

  StableSortExpander stabilizer;
  EXPECT_TRUE(stabilizer.Run(module.get()).ValueOrDie());
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, GmockMatch(m::GetTupleElement(
                        m::Sort(m::Parameter(0), m::Iota()), 0)));
  CheckComputationHasTieBreaker(
      root->operand(0)->to_apply()->root_instruction(), /*iota_parameter=*/1);
}

TEST_F(StableSortExpanderTest,
       StabilizeSortReuseIotaOperandComplicatedComparison) {