(i)
| 1016 | zeros = array_ops.zeros([state_size]) |
| 1017 | |
| 1018 | def loop_fn(i): |
| 1019 | sequence_length_i = array_ops.gather(sequence_length, i) |
| 1020 | |
| 1021 | def body_fn(t, state, ta): |
| 1022 | inputs_t = array_ops.expand_dims( |
| 1023 | array_ops.gather(inputs_ta.read(t), i), 0) |
| 1024 | output, new_state = cell(inputs_t, state) |
| 1025 | output = array_ops.reshape(output, [-1]) |
| 1026 | # TODO(agarwal): one optimization that dynamic_rnn uses is to avoid the |
| 1027 | # array_ops.where when t < min(sequence_length). Doing that requires |
| 1028 | # supporting tf.cond pfor conversion. |
| 1029 | done = t >= sequence_length_i |
| 1030 | output = array_ops.where(done, zeros, output) |
| 1031 | ta = ta.write(t, output) |
| 1032 | new_state = [array_ops.where(done, s, ns) for s, ns in |
| 1033 | zip(nest.flatten(state), nest.flatten(new_state))] |
| 1034 | new_state = nest.pack_sequence_as(state, new_state) |
| 1035 | return t + 1, new_state, ta |
| 1036 | |
| 1037 | def condition_fn(t, _, unused): |
| 1038 | del unused |
| 1039 | return t < max_steps |
| 1040 | |
| 1041 | initial_state = cell.zero_state(1, dtypes.float32) |
| 1042 | _, state, ta = control_flow_ops.while_loop(condition_fn, body_fn, [ |
| 1043 | 0, initial_state, |
| 1044 | tensor_array_ops.TensorArray(dtypes.float32, max_steps) |
| 1045 | ]) |
| 1046 | |
| 1047 | new_state = [array_ops.reshape(x, [-1]) for x in nest.flatten(state)] |
| 1048 | new_state = nest.pack_sequence_as(initial_state, new_state) |
| 1049 | return ta.stack(), new_state |
| 1050 | |
| 1051 | pfor_output = pfor_control_flow_ops.pfor(loop_fn, batch_size) |
| 1052 | tf_output = rnn.dynamic_rnn( |
nothing calls this directly
no test coverage detected