Input object passed to registered pfor converters.
| 665 | |
| 666 | |
| 667 | class _PforInput(object): |
| 668 | """Input object passed to registered pfor converters.""" |
| 669 | |
| 670 | def __init__(self, pfor, op, inputs): |
| 671 | """Creates a _PforInput object. |
| 672 | |
| 673 | Args: |
| 674 | pfor: PFor converter object. |
| 675 | op: the Operation object that is being converted. |
| 676 | inputs: list of WrappedTensor objects representing converted values of the |
| 677 | inputs of `op`. |
| 678 | """ |
| 679 | self.pfor = pfor |
| 680 | self._op = op |
| 681 | self._inputs = inputs |
| 682 | |
| 683 | def stack_inputs(self, stack_indices=None): |
| 684 | """Stacks unstacked inputs at `stack_indices`. |
| 685 | |
| 686 | Args: |
| 687 | stack_indices: indices of inputs at which stacking is done. If None, |
| 688 | stacking is done at all indices. |
| 689 | """ |
| 690 | if stack_indices is None: |
| 691 | stack_indices = range(len(self._inputs)) |
| 692 | length = self.pfor.loop_len_vector |
| 693 | for i in stack_indices: |
| 694 | inp = self._inputs[i] |
| 695 | if not inp.is_stacked: |
| 696 | self._inputs[i] = _stack(inp.t, length) |
| 697 | |
| 698 | def expanddim_inputs_for_broadcast(self): |
| 699 | """Reshapes stacked inputs to prepare them for broadcast. |
| 700 | |
| 701 | Since stacked inputs have an extra leading dimension, automatic broadcasting |
| 702 | rules could incorrectly try to expand dimensions before that leading |
| 703 | dimension. To avoid that, we reshape these stacked inputs to the maximum |
| 704 | rank they will need to be broadcasted to. |
| 705 | """ |
| 706 | if not self._inputs: |
| 707 | return |
| 708 | |
| 709 | # Find max rank |
| 710 | def _get_rank(x): |
| 711 | rank = array_ops.rank(x.t) |
| 712 | if not x.is_stacked: |
| 713 | rank += 1 |
| 714 | return rank |
| 715 | |
| 716 | ranks = [_get_rank(x) for x in self._inputs] |
| 717 | max_rank = ranks[0] |
| 718 | for rank in ranks[1:]: |
| 719 | max_rank = math_ops.maximum(rank, max_rank) |
| 720 | |
| 721 | for i, inp in enumerate(self._inputs): |
| 722 | if inp.is_stacked: |
| 723 | shape = array_ops.shape(inp.t) |
| 724 | rank_diff = array_ops.reshape(max_rank - ranks[i], [1]) |