| 1072 | } |
| 1073 | |
| 1074 | Status ShapeVerifier::CheckShape(const HloInstruction* instruction, |
| 1075 | const Shape& inferred_shape, |
| 1076 | bool only_compare_minor_to_major_in_layout) { |
| 1077 | // If allow_mixed_precision_ is false, check if there are operands with |
| 1078 | // different precisions. We need this check because ShapeInference allows |
| 1079 | // mixed precision inputs. |
| 1080 | if (!allow_mixed_precision_) { |
| 1081 | TF_RETURN_IF_ERROR(CheckMixedPrecisionOperands(instruction)); |
| 1082 | } |
| 1083 | |
| 1084 | // Check if the output shape matches the expected shape. |
| 1085 | // |
| 1086 | // We treat BF16 and F32 as compatible types if mixed precision is allowed, |
| 1087 | // but only when the instruction defines the BF16/F32 buffer. |
| 1088 | bool equal = [&] { |
| 1089 | switch (instruction->opcode()) { |
| 1090 | // The opcodes below can't have implicit layout conversions, nor can they |
| 1091 | // implicitly transform f32 -> bf16. Fundamentally these are either |
| 1092 | // reinterpreting existing data (e.g. kBitcast) or shuffling data around |
| 1093 | // without modifying it (e.g. kGetTupleElement, kTupleSelect). |
| 1094 | case HloOpcode::kBitcast: |
| 1095 | case HloOpcode::kCall: |
| 1096 | case HloOpcode::kConditional: |
| 1097 | case HloOpcode::kConstant: |
| 1098 | case HloOpcode::kCopyDone: |
| 1099 | case HloOpcode::kCopyStart: |
| 1100 | case HloOpcode::kCustomCall: |
| 1101 | case HloOpcode::kGetTupleElement: |
| 1102 | case HloOpcode::kInfeed: |
| 1103 | case HloOpcode::kOutfeed: |
| 1104 | case HloOpcode::kParameter: |
| 1105 | case HloOpcode::kRecv: |
| 1106 | case HloOpcode::kRecvDone: |
| 1107 | case HloOpcode::kSend: |
| 1108 | case HloOpcode::kSendDone: |
| 1109 | case HloOpcode::kTuple: |
| 1110 | case HloOpcode::kTupleSelect: |
| 1111 | case HloOpcode::kWhile: |
| 1112 | return ShapesSame(instruction->shape(), inferred_shape, |
| 1113 | only_compare_minor_to_major_in_layout); |
| 1114 | |
| 1115 | // We allow arbitrary layout and f32->bf16 transformations on all other |
| 1116 | // instructions, although this may be made more strict pending discussion |
| 1117 | // in b/112709536. |
| 1118 | default: |
| 1119 | if (allow_mixed_precision_) { |
| 1120 | return ShapeUtil::CompatibleIgnoringFpPrecision(instruction->shape(), |
| 1121 | inferred_shape); |
| 1122 | } else { |
| 1123 | return ShapeUtil::Compatible(instruction->shape(), inferred_shape); |
| 1124 | } |
| 1125 | } |
| 1126 | }(); |
| 1127 | if (!equal) { |
| 1128 | return InternalError( |
| 1129 | "Expected instruction to have shape equal to %s, actual " |
| 1130 | "shape is %s:\n%s", |
| 1131 | StringifyShape(inferred_shape), StringifyShape(instruction->shape()), |
nothing calls this directly
no test coverage detected