| 268 | } |
| 269 | |
| 270 | XLA_TEST_F(PadTest, Pad4DU8Array) { |
| 271 | XlaBuilder b(TestName()); |
| 272 | auto input = absl::make_unique<Array4D<uint8>>(1, 1, 3, 2); |
| 273 | Array2D<uint8> input_xy({ |
| 274 | {1, 2}, // row 0 |
| 275 | {3, 4}, // row 1 |
| 276 | {5, 6}, // row 2 |
| 277 | }); |
| 278 | input->FillWithYX(input_xy); |
| 279 | |
| 280 | Pad(AddParam(*input, &b), ConstantR0<uint8>(&b, 35), |
| 281 | r4_padding_on_dim0_dim1_); |
| 282 | |
| 283 | auto expected = absl::make_unique<Array4D<uint8>>(2, 3, 3, 2); |
| 284 | expected->Fill(35); |
| 285 | (*expected)(1, 0, 0, 0) = 1; |
| 286 | (*expected)(1, 0, 0, 1) = 2; |
| 287 | (*expected)(1, 0, 1, 0) = 3; |
| 288 | (*expected)(1, 0, 1, 1) = 4; |
| 289 | (*expected)(1, 0, 2, 0) = 5; |
| 290 | (*expected)(1, 0, 2, 1) = 6; |
| 291 | ComputeAndCompareR4<uint8>(&b, *expected, {}); |
| 292 | } |
| 293 | |
| 294 | XLA_TEST_F(PadTest, Pad4DPredArray) { |
| 295 | XlaBuilder b(TestName()); |