| 28 | using ArithmeticTest = ClientLibraryTestBase; |
| 29 | |
| 30 | XLA_TEST_F(ArithmeticTest, ArgMinR2Axis0) { |
| 31 | XlaBuilder builder(TestName()); |
| 32 | auto x = ConstantR2<int32>(&builder, {{1, 7, 4}, {6, 3, 5}, {8, 3, 3}}); |
| 33 | ArgMin(x, S32, /*axis=*/0); |
| 34 | |
| 35 | std::vector<int32> expected = {0, 2, 2}; |
| 36 | ComputeAndCompareR1<int32>(&builder, expected, {}); |
| 37 | } |
| 38 | |
| 39 | XLA_TEST_F(ArithmeticTest, ArgMinR2Axis1) { |
| 40 | XlaBuilder builder(TestName()); |