| 79 | }; |
| 80 | |
| 81 | TEST_P(FloatReverseTest, Reverses) { |
| 82 | const ReverseSpec& spec = GetParam(); |
| 83 | std::vector<float> input_vector( |
| 84 | ShapeUtil::ElementsIn(ShapeUtil::MakeShape(F32, spec.input_dims))); |
| 85 | std::iota(input_vector.begin(), input_vector.end(), 0.0); |
| 86 | auto r1_literal = LiteralUtil::CreateR1<float>(input_vector); |
| 87 | auto input_literal = r1_literal.Reshape(spec.input_dims).ConsumeValueOrDie(); |
| 88 | |
| 89 | XlaBuilder builder(TestName()); |
| 90 | auto a = AddParam(input_literal, &builder); |
| 91 | Rev(a, spec.reversal); |
| 92 | |
| 93 | Literal expected = input_literal.Clone(); |
| 94 | std::vector<int64> output_indices(spec.input_dims.size()); |
| 95 | expected.EachCell<float>([&](absl::Span<const int64> indices, float) { |
| 96 | for (int64 i = 0; i < indices.size(); ++i) { |
| 97 | output_indices[i] = indices[i]; |
| 98 | } |
| 99 | float value = input_literal.Get<float>(indices); |
| 100 | for (int64 dim : spec.reversal) { |
| 101 | output_indices[dim] = (spec.input_dims[dim] - 1) - indices[dim]; |
| 102 | } |
| 103 | expected.Set<float>(output_indices, value); |
| 104 | }); |
| 105 | ComputeAndCompareLiteral(&builder, expected, {}); |
| 106 | } |
| 107 | |
| 108 | INSTANTIATE_TEST_CASE_P(FloatReverseInstance, FloatReverseTest, |
| 109 | ::testing::ValuesIn(GetTestCases()), |