| 756 | } |
| 757 | |
| 758 | XlaOp XlaBuilder::BroadcastInDim( |
| 759 | XlaOp operand, const absl::Span<const int64> out_dim_size, |
| 760 | const absl::Span<const int64> broadcast_dimensions) { |
| 761 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 762 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 763 | // Output shape, in the case of degenerate broadcast, the out_dim_size is |
| 764 | // not necessarily the same as the dimension sizes of the output shape. |
| 765 | auto output_shape = |
| 766 | ShapeUtil::MakeShape(operand_shape->element_type(), out_dim_size); |
| 767 | if (operand_shape->rank() != broadcast_dimensions.size()) { |
| 768 | return InvalidArgument( |
| 769 | "Size of broadcast_dimensions has to match operand's rank; operand " |
| 770 | "rank: %lld, size of broadcast_dimensions %u.", |
| 771 | operand_shape->rank(), broadcast_dimensions.size()); |
| 772 | } |
| 773 | for (int i = 0; i < broadcast_dimensions.size(); i++) { |
| 774 | if (broadcast_dimensions[i] < 0 || |
| 775 | broadcast_dimensions[i] > out_dim_size.size()) { |
| 776 | return InvalidArgument("Broadcast dimension %lld is out of bound", |
| 777 | broadcast_dimensions[i]); |
| 778 | } |
| 779 | output_shape.set_dynamic_dimension( |
| 780 | broadcast_dimensions[i], operand_shape->is_dynamic_dimension(i)); |
| 781 | } |
| 782 | |
| 783 | TF_RETURN_IF_ERROR(ShapeInference::InferBroadcastShape( |
| 784 | *operand_shape, output_shape, broadcast_dimensions) |
| 785 | .status()); |
| 786 | std::vector<int64> in_dim_size(out_dim_size.begin(), out_dim_size.end()); |
| 787 | for (int i = 0; i < broadcast_dimensions.size(); i++) { |
| 788 | in_dim_size[broadcast_dimensions[i]] = operand_shape->dimensions(i); |
| 789 | } |
| 790 | const auto& in_dim_shape = |
| 791 | ShapeUtil::MakeShape(operand_shape->element_type(), in_dim_size); |
| 792 | TF_ASSIGN_OR_RETURN( |
| 793 | XlaOp in_dim_broadcast, |
| 794 | InDimBroadcast(in_dim_shape, operand, broadcast_dimensions)); |
| 795 | |
| 796 | // If broadcast is not degenerate, return broadcasted result. |
| 797 | if (ShapeUtil::Equal(in_dim_shape, output_shape)) { |
| 798 | return in_dim_broadcast; |
| 799 | } |
| 800 | |
| 801 | // Otherwise handle degenerate broadcast case. |
| 802 | return AddBroadcastSequence(output_shape, in_dim_broadcast); |
| 803 | }); |
| 804 | } |
| 805 | |
| 806 | StatusOr<XlaOp> XlaBuilder::ReshapeInternal(const Shape& shape, XlaOp operand, |
| 807 | int64 inferred_dimension) { |