| 758 | } |
| 759 | |
| 760 | Status ShapeVerifier::HandleMap(HloInstruction* map) { |
| 761 | std::vector<const Shape*> operand_shapes; |
| 762 | int64 max_operand_rank = 0; |
| 763 | for (const HloInstruction* operand : map->operands()) { |
| 764 | operand_shapes.push_back(&operand->shape()); |
| 765 | max_operand_rank = std::max(max_operand_rank, operand->shape().rank()); |
| 766 | } |
| 767 | // TODO(b/65689298) Remove code below once Map is generalized to accept |
| 768 | // arbitrary map dimensions. |
| 769 | std::vector<int64> map_dims(max_operand_rank); |
| 770 | std::iota(map_dims.begin(), map_dims.end(), 0); |
| 771 | |
| 772 | TF_RETURN_IF_ERROR(CheckShape( |
| 773 | map, |
| 774 | ShapeInference::InferMapShape( |
| 775 | operand_shapes, map->to_apply()->ComputeProgramShape(), map_dims))); |
| 776 | |
| 777 | return allow_mixed_precision_ |
| 778 | ? Status::OK() |
| 779 | : SameElementTypesForOperandsAndToApplyParameters( |
| 780 | *map, map->operands().size()); |
| 781 | } |
| 782 | |
| 783 | Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) { |
| 784 | TF_RETURN_IF_ERROR(CheckShape( |
nothing calls this directly
no test coverage detected