Calculate one step of a dynamic RNN minibatch. Returns an (output, state) pair conditioned on the sequence_lengths. When skip_conditionals=False, the pseudocode is something like: if t >= max_sequence_length: return (zero_output, state) if t < min_sequence_length: return call_cell(
(
time, sequence_length, min_sequence_length, max_sequence_length,
zero_output, state, call_cell, state_size, skip_conditionals=False)
| 136 | |
| 137 | # pylint: disable=unused-argument |
| 138 | def _rnn_step( |
| 139 | time, sequence_length, min_sequence_length, max_sequence_length, |
| 140 | zero_output, state, call_cell, state_size, skip_conditionals=False): |
| 141 | """Calculate one step of a dynamic RNN minibatch. |
| 142 | |
| 143 | Returns an (output, state) pair conditioned on the sequence_lengths. |
| 144 | When skip_conditionals=False, the pseudocode is something like: |
| 145 | |
| 146 | if t >= max_sequence_length: |
| 147 | return (zero_output, state) |
| 148 | if t < min_sequence_length: |
| 149 | return call_cell() |
| 150 | |
| 151 | # Selectively output zeros or output, old state or new state depending |
| 152 | # on if we've finished calculating each row. |
| 153 | new_output, new_state = call_cell() |
| 154 | final_output = np.vstack([ |
| 155 | zero_output if time >= sequence_lengths[r] else new_output_r |
| 156 | for r, new_output_r in enumerate(new_output) |
| 157 | ]) |
| 158 | final_state = np.vstack([ |
| 159 | state[r] if time >= sequence_lengths[r] else new_state_r |
| 160 | for r, new_state_r in enumerate(new_state) |
| 161 | ]) |
| 162 | return (final_output, final_state) |
| 163 | |
| 164 | Args: |
| 165 | time: Python int, the current time step |
| 166 | sequence_length: int32 `Tensor` vector of size [batch_size] |
| 167 | min_sequence_length: int32 `Tensor` scalar, min of sequence_length |
| 168 | max_sequence_length: int32 `Tensor` scalar, max of sequence_length |
| 169 | zero_output: `Tensor` vector of shape [output_size] |
| 170 | state: Either a single `Tensor` matrix of shape `[batch_size, state_size]`, |
| 171 | or a list/tuple of such tensors. |
| 172 | call_cell: lambda returning tuple of (new_output, new_state) where |
| 173 | new_output is a `Tensor` matrix of shape `[batch_size, output_size]`. |
| 174 | new_state is a `Tensor` matrix of shape `[batch_size, state_size]`. |
| 175 | state_size: The `cell.state_size` associated with the state. |
| 176 | skip_conditionals: Python bool, whether to skip using the conditional |
| 177 | calculations. This is useful for `dynamic_rnn`, where the input tensor |
| 178 | matches `max_sequence_length`, and using conditionals just slows |
| 179 | everything down. |
| 180 | |
| 181 | Returns: |
| 182 | A tuple of (`final_output`, `final_state`) as given by the pseudocode above: |
| 183 | final_output is a `Tensor` matrix of shape [batch_size, output_size] |
| 184 | final_state is either a single `Tensor` matrix, or a tuple of such |
| 185 | matrices (matching length and shapes of input `state`). |
| 186 | |
| 187 | Raises: |
| 188 | ValueError: If the cell returns a state tuple whose length does not match |
| 189 | that returned by `state_size`. |
| 190 | """ |
| 191 | |
| 192 | # Convert state to a list for ease of use |
| 193 | flat_state = nest.flatten(state) |
| 194 | flat_zero_output = nest.flatten(zero_output) |
| 195 |
no test coverage detected