| 805 | } |
| 806 | |
| 807 | Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { |
| 808 | TF_RETURN_IF_ERROR( |
| 809 | CheckParameterCount(xla_while, xla_while->while_body(), 1)); |
| 810 | TF_RETURN_IF_ERROR( |
| 811 | CheckParameterCount(xla_while, xla_while->while_condition(), 1)); |
| 812 | TF_RETURN_IF_ERROR( |
| 813 | CheckOperandAndParameter(xla_while, 0, xla_while->while_body(), 0)); |
| 814 | TF_RETURN_IF_ERROR( |
| 815 | CheckOperandAndParameter(xla_while, 0, xla_while->while_condition(), 0)); |
| 816 | const Shape& conditional_shape = |
| 817 | xla_while->while_condition()->root_instruction()->shape(); |
| 818 | if (!ShapeUtil::Compatible(conditional_shape, |
| 819 | ShapeUtil::MakeShape(PRED, {}))) { |
| 820 | return InternalError( |
| 821 | "Conditional computation shape does not lead to a scalar predicate " |
| 822 | "shape: %s", |
| 823 | StringifyShape(conditional_shape)); |
| 824 | } |
| 825 | // The shape of kWhile should match the shape of the body computation it |
| 826 | // calls. |
| 827 | return CheckShape(xla_while, |
| 828 | xla_while->while_body()->root_instruction()->shape()); |
| 829 | } |
| 830 | |
| 831 | Status ShapeVerifier::HandleConditional(HloInstruction* conditional) { |
| 832 | if (!ShapeUtil::IsScalar(conditional->operand(0)->shape())) { |
no test coverage detected