For a while instruction, all the following layouts must be the same: (1) init operand (2) condition computation parameter (3) body computation parameter (4) body computation result (5) while instruction result
| 710 | // (4) body computation result |
| 711 | // (5) while instruction result |
| 712 | Status CheckWhileLayout(HloInstruction* while_inst, |
| 713 | const ComputationLayout& condition_computation_layout, |
| 714 | const ComputationLayout& body_computation_layout) { |
| 715 | auto init_shape = while_inst->operand(0)->shape(); |
| 716 | TF_RET_CHECK( |
| 717 | condition_computation_layout.parameter_layout(0).MatchesLayoutInShape( |
| 718 | init_shape, /*minor_to_major_only=*/true)); |
| 719 | TF_RET_CHECK(body_computation_layout.parameter_layout(0).MatchesLayoutInShape( |
| 720 | init_shape, /*minor_to_major_only=*/true)); |
| 721 | TF_RET_CHECK(body_computation_layout.result_layout().MatchesLayoutInShape( |
| 722 | init_shape, /*minor_to_major_only=*/true)); |
| 723 | TF_RET_CHECK(LayoutsInShapesEqual(init_shape, while_inst->shape())); |
| 724 | return Status::OK(); |
| 725 | } |
| 726 | |
| 727 | Status CheckConditionalLayout( |
| 728 | HloInstruction* instruction, |
no test coverage detected