| 240 | }; |
| 241 | |
| 242 | XLA_TEST_F(ExecutionTest, ScatterUpdate) { |
| 243 | // Test that scattering on indices=[2] is same as scattering on indices=[4] |
| 244 | // and dynamic dimension = 2 |
| 245 | const string hlo_text = R"( |
| 246 | HloModule TensorFlowScatterV1 |
| 247 | |
| 248 | update_s32 (lhs: s32[], rhs: s32[]) -> s32[] { |
| 249 | lhs = s32[] parameter(0) |
| 250 | ROOT rhs = s32[] parameter(1) |
| 251 | } |
| 252 | |
| 253 | ENTRY main { |
| 254 | operand = s32[3,3] parameter(0) |
| 255 | indices = s32[INDICES_BOUND] parameter(1) |
| 256 | updates = s32[INDICES_BOUND,3] parameter(2) |
| 257 | dynamic_size = s32[] parameter(3) |
| 258 | ROOT scatter = s32[3,3] scatter(operand, indices, updates), |
| 259 | to_apply=update_s32, |
| 260 | update_window_dims={1}, |
| 261 | inserted_window_dims={0}, |
| 262 | scatter_dims_to_operand_dims={0}, |
| 263 | index_vector_dim=1 |
| 264 | |
| 265 | } |
| 266 | )"; |
| 267 | const string hlo_text_not_padded = |
| 268 | absl::StrReplaceAll(hlo_text, {{"INDICES_BOUND", "2"}}); |
| 269 | auto module_not_padded = GetHloModule(hlo_text_not_padded); |
| 270 | |
| 271 | Literal operand = |
| 272 | LiteralUtil::CreateR2<int32>({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); |
| 273 | Literal scatter_indices = LiteralUtil::CreateR1<int32>({0, 2}); |
| 274 | Literal updates = LiteralUtil::CreateR2<int32>({{10, 20, 30}, {70, 80, 90}}); |
| 275 | Literal dynamic_size = LiteralUtil::CreateR0<int32>(2); |
| 276 | |
| 277 | Literal not_padded = |
| 278 | ExecuteAndTransfer(std::move(module_not_padded), |
| 279 | {&operand, &scatter_indices, &updates, &dynamic_size}); |
| 280 | |
| 281 | // Pad input to 4. |
| 282 | const string hlo_text_padded = |
| 283 | absl::StrReplaceAll(hlo_text, {{"INDICES_BOUND", "4"}}); |
| 284 | auto module_padded = GetHloModule(hlo_text_padded); |
| 285 | // Set up dynamic parameter binding. |
| 286 | TF_CHECK_OK(module_padded->dynamic_parameter_binding().Bind( |
| 287 | DynamicParameterBinding::DynamicParameter{3, {}}, |
| 288 | DynamicParameterBinding::DynamicDimension{1, {}, 0})); |
| 289 | TF_CHECK_OK(module_padded->dynamic_parameter_binding().Bind( |
| 290 | DynamicParameterBinding::DynamicParameter{3, {}}, |
| 291 | DynamicParameterBinding::DynamicDimension{2, {}, 0})); |
| 292 | // Pad the rest of input with garbage data. |
| 293 | Literal scatter_indices_padded = LiteralUtil::CreateR1<int32>({0, 2, 0, 4}); |
| 294 | Literal updates_padded = LiteralUtil::CreateR2<int32>( |
| 295 | {{10, 20, 30}, {70, 80, 90}, {30, 22, 11}, {-1, 20, -1}}); |
| 296 | DynamicPadder padder; |
| 297 | TF_CHECK_OK(padder.Run(module_padded.get()).status()); |
| 298 | Literal padded = PadAndExecute( |
| 299 | std::move(module_padded), |
nothing calls this directly
no test coverage detected