Take a time step of the dynamic RNN. Args: time: int32 scalar Tensor. output_ta_t: List of `TensorArray`s that represent the output. state: nested tuple of vector tensors that represent the state. Returns: The tuple (time + 1, output_ta_t with updated flow, new_stat
(time, output_ta_t, state)
| 843 | input_ta = flat_input |
| 844 | |
| 845 | def _time_step(time, output_ta_t, state): |
| 846 | """Take a time step of the dynamic RNN. |
| 847 | |
| 848 | Args: |
| 849 | time: int32 scalar Tensor. |
| 850 | output_ta_t: List of `TensorArray`s that represent the output. |
| 851 | state: nested tuple of vector tensors that represent the state. |
| 852 | |
| 853 | Returns: |
| 854 | The tuple (time + 1, output_ta_t with updated flow, new_state). |
| 855 | """ |
| 856 | |
| 857 | if in_graph_mode: |
| 858 | input_t = tuple(ta.read(time) for ta in input_ta) |
| 859 | # Restore some shape information |
| 860 | for input_, shape in zip(input_t, inputs_got_shape): |
| 861 | input_.set_shape(shape[1:]) |
| 862 | else: |
| 863 | input_t = tuple(ta[time.numpy()] for ta in input_ta) |
| 864 | |
| 865 | input_t = nest.pack_sequence_as(structure=inputs, flat_sequence=input_t) |
| 866 | # Keras RNN cells only accept state as list, even if it's a single tensor. |
| 867 | is_keras_rnn_cell = _is_keras_rnn_cell(cell) |
| 868 | if is_keras_rnn_cell and not nest.is_sequence(state): |
| 869 | state = [state] |
| 870 | call_cell = lambda: cell(input_t, state) |
| 871 | |
| 872 | if sequence_length is not None: |
| 873 | (output, new_state) = _rnn_step( |
| 874 | time=time, |
| 875 | sequence_length=sequence_length, |
| 876 | min_sequence_length=min_sequence_length, |
| 877 | max_sequence_length=max_sequence_length, |
| 878 | zero_output=zero_output, |
| 879 | state=state, |
| 880 | call_cell=call_cell, |
| 881 | state_size=state_size, |
| 882 | skip_conditionals=True) |
| 883 | else: |
| 884 | (output, new_state) = call_cell() |
| 885 | |
| 886 | # Keras cells always wrap state as list, even if it's a single tensor. |
| 887 | if is_keras_rnn_cell and len(new_state) == 1: |
| 888 | new_state = new_state[0] |
| 889 | # Pack state if using state tuples |
| 890 | output = nest.flatten(output) |
| 891 | |
| 892 | if in_graph_mode: |
| 893 | output_ta_t = tuple( |
| 894 | ta.write(time, out) for ta, out in zip(output_ta_t, output)) |
| 895 | else: |
| 896 | for ta, out in zip(output_ta_t, output): |
| 897 | ta[time.numpy()] = out |
| 898 | |
| 899 | return (time + 1, output_ta_t, new_state) |
| 900 | |
| 901 | if in_graph_mode: |
| 902 | # Make sure that we run at least 1 step, if necessary, to ensure |