| 35 | public ::testing::WithParamInterface<std::tuple<PrimitiveType, int>> {}; |
| 36 | |
| 37 | XLA_TEST_P(IotaR1Test, DoIt) { |
| 38 | const auto& spec = GetParam(); |
| 39 | const auto element_type = std::get<0>(spec); |
| 40 | const int64 num_elements = std::get<1>(spec); |
| 41 | XlaBuilder builder(TestName() + "_" + PrimitiveType_Name(element_type)); |
| 42 | Iota(&builder, element_type, num_elements); |
| 43 | if (element_type == F32) { |
| 44 | ComputeAndCompareR1<float>(&builder, GetR1Expected<float>(num_elements), {}, |
| 45 | ErrorSpec{0.0001}); |
| 46 | } else if (element_type == U32) { |
| 47 | ComputeAndCompareR1<uint32>(&builder, GetR1Expected<uint32>(num_elements), |
| 48 | {}); |
| 49 | } else { |
| 50 | CHECK_EQ(element_type, S32); |
| 51 | ComputeAndCompareR1<int32>(&builder, GetR1Expected<int32>(num_elements), |
| 52 | {}); |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | INSTANTIATE_TEST_CASE_P(IotaR1TestInstantiation, IotaR1Test, |
| 57 | ::testing::Combine(::testing::Values(F32, U32, S32), |