| 35 | class ScatterExpanderTest : public HloTestBase {}; |
| 36 | |
| 37 | TEST_F(ScatterExpanderTest, ScatterOperandWithoutLayout) { |
| 38 | const char* kModuleStr = R"( |
| 39 | HloModule scatter_expander |
| 40 | |
| 41 | scatter_computation { |
| 42 | parameter0 = s32[] parameter(0) |
| 43 | ROOT parameter1 = s32[] parameter(1) |
| 44 | } |
| 45 | |
| 46 | ENTRY kernel_entry { |
| 47 | operand = s32[5] iota(), iota_dimension=0 |
| 48 | indices = s32[1] parameter(0) |
| 49 | update = s32[] constant(0) |
| 50 | ROOT scatter = s32[5]{0} scatter(operand, indices, update), |
| 51 | update_window_dims={}, inserted_window_dims={0}, |
| 52 | scatter_dims_to_operand_dims={0}, index_vector_dim=0, |
| 53 | to_apply=scatter_computation |
| 54 | })"; |
| 55 | |
| 56 | TF_ASSERT_OK_AND_ASSIGN(auto module, |
| 57 | ParseAndReturnVerifiedModule(kModuleStr)); |
| 58 | |
| 59 | // The HLO parser changes all no layout shapes from the input to have a |
| 60 | // default layout, clear the layout of the scatter operand for testing. |
| 61 | HloInstruction* scatter_operand = FindInstruction(module.get(), "operand"); |
| 62 | scatter_operand->mutable_shape()->clear_layout(); |
| 63 | |
| 64 | ScatterExpander scatter_expander; |
| 65 | TF_ASSERT_OK_AND_ASSIGN(bool result, |
| 66 | RunHloPass(&scatter_expander, module.get())); |
| 67 | EXPECT_TRUE(result); |
| 68 | } |
| 69 | |
| 70 | } // namespace |
| 71 | } // namespace xla |
nothing calls this directly
no test coverage detected