Helper functions for is_broadcast_compatible and broadcast_shape. Args: shape_x: A `TensorShape` shape_y: A `TensorShape` Returns: Returns None if the shapes are not broadcast compatible, a list of the broadcast dimensions otherwise.
(shape_x, shape_y)
| 512 | |
| 513 | |
| 514 | def _broadcast_shape_helper(shape_x, shape_y): |
| 515 | """Helper functions for is_broadcast_compatible and broadcast_shape. |
| 516 | |
| 517 | Args: |
| 518 | shape_x: A `TensorShape` |
| 519 | shape_y: A `TensorShape` |
| 520 | |
| 521 | Returns: |
| 522 | Returns None if the shapes are not broadcast compatible, |
| 523 | a list of the broadcast dimensions otherwise. |
| 524 | """ |
| 525 | # To compute the broadcasted dimensions, we zip together shape_x and shape_y, |
| 526 | # and pad with 1 to make them the same length. |
| 527 | broadcasted_dims = reversed(list(six.moves.zip_longest( |
| 528 | reversed(shape_x.dims), |
| 529 | reversed(shape_y.dims), |
| 530 | fillvalue=tensor_shape.Dimension(1)))) |
| 531 | # Next we combine the dimensions according to the numpy broadcasting rules. |
| 532 | # http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html |
| 533 | return_dims = [] |
| 534 | for (dim_x, dim_y) in broadcasted_dims: |
| 535 | if dim_x.value is None or dim_y.value is None: |
| 536 | # One or both dimensions is unknown. If either dimension is greater than |
| 537 | # 1, we assume that the program is correct, and the other dimension will |
| 538 | # be broadcast to match it. |
| 539 | # TODO(mrry): If we eliminate the shape checks in C++, we must still |
| 540 | # assert that the unknown dim is either 1 or the same as the known dim. |
| 541 | if dim_x.value is not None and dim_x.value > 1: |
| 542 | return_dims.append(dim_x) |
| 543 | elif dim_y.value is not None and dim_y.value > 1: |
| 544 | return_dims.append(dim_y) |
| 545 | else: |
| 546 | return_dims.append(None) |
| 547 | elif dim_x.value == 1: |
| 548 | # We will broadcast dim_x to dim_y. |
| 549 | return_dims.append(dim_y) |
| 550 | elif dim_y.value == 1: |
| 551 | # We will broadcast dim_y to dim_x. |
| 552 | return_dims.append(dim_x) |
| 553 | elif dim_x.value == dim_y.value: |
| 554 | # The dimensions are compatible, so output is the same size in that |
| 555 | # dimension. |
| 556 | return_dims.append(dim_x.merge_with(dim_y)) |
| 557 | else: |
| 558 | return None |
| 559 | return return_dims |
| 560 | |
| 561 | |
| 562 | def is_broadcast_compatible(shape_x, shape_y): |
no test coverage detected