RNN step function. Arguments: time: Current timestep value. output_ta_t: TensorArray. prev_output: tuple of outputs from time - 1. *states: List of states. Returns: Tuple: `(time + 1, output_ta_t, output) + tuple(new_state
(time, output_ta_t, prev_output, *states)
| 3968 | flat_zero_output = tuple(array_ops.zeros_like(o) |
| 3969 | for o in nest.flatten(output_time_zero)) |
| 3970 | def _step(time, output_ta_t, prev_output, *states): |
| 3971 | """RNN step function. |
| 3972 | |
| 3973 | Arguments: |
| 3974 | time: Current timestep value. |
| 3975 | output_ta_t: TensorArray. |
| 3976 | prev_output: tuple of outputs from time - 1. |
| 3977 | *states: List of states. |
| 3978 | |
| 3979 | Returns: |
| 3980 | Tuple: `(time + 1, output_ta_t, output) + tuple(new_states)` |
| 3981 | """ |
| 3982 | current_input = tuple(ta.read(time) for ta in input_ta) |
| 3983 | # maybe set shape. |
| 3984 | current_input = nest.pack_sequence_as(inputs, current_input) |
| 3985 | mask_t = mask_ta.read(time) |
| 3986 | output, new_states = step_function(current_input, |
| 3987 | tuple(states) + tuple(constants)) |
| 3988 | # mask output |
| 3989 | flat_output = nest.flatten(output) |
| 3990 | flat_mask_output = (flat_zero_output if zero_output_for_mask |
| 3991 | else nest.flatten(prev_output)) |
| 3992 | tiled_mask_t = tuple(_expand_mask(mask_t, o) for o in flat_output) |
| 3993 | flat_new_output = tuple( |
| 3994 | array_ops.where(m, o, zo) for m, o, zo in zip( |
| 3995 | tiled_mask_t, flat_output, flat_mask_output)) |
| 3996 | |
| 3997 | # mask states |
| 3998 | flat_state = nest.flatten(states) |
| 3999 | flat_new_state = nest.flatten(new_states) |
| 4000 | for state, new_state in zip(flat_state, flat_new_state): |
| 4001 | if isinstance(new_state, ops.Tensor): |
| 4002 | new_state.set_shape(state.shape) |
| 4003 | tiled_mask_t = tuple(_expand_mask(mask_t, s) for s in flat_state) |
| 4004 | flat_final_state = tuple( |
| 4005 | array_ops.where(m, s, ps) |
| 4006 | for m, s, ps in zip(tiled_mask_t, flat_new_state, flat_state)) |
| 4007 | new_states = nest.pack_sequence_as(new_states, flat_final_state) |
| 4008 | |
| 4009 | output_ta_t = tuple( |
| 4010 | ta.write(time, out) |
| 4011 | for ta, out in zip(output_ta_t, flat_new_output)) |
| 4012 | return (time + 1, output_ta_t, |
| 4013 | tuple(flat_new_output)) + tuple(new_states) |
| 4014 | |
| 4015 | final_outputs = control_flow_ops.while_loop( |
| 4016 | body=_step, |