MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / body

Function body

modelzoo/features/multihash_variable/dien/script/rnn.py:1044–1093  ·  view source on GitHub ↗

Internal while loop body for raw_rnn. Args: time: time scalar. elements_finished: batch-size vector. current_input: possibly nested tuple of input tensors. emit_ta: possibly nested tuple of output TensorArrays. state: possibly nested tuple of state tens

(time, elements_finished, current_input,
             emit_ta, state, loop_state)

Source from the content-addressed store, hash-verified

1042 return math_ops.logical_not(math_ops.reduce_all(elements_finished))
1043
1044 def body(time, elements_finished, current_input,
1045 emit_ta, state, loop_state):
1046 """Internal while loop body for raw_rnn.
1047
1048 Args:
1049 time: time scalar.
1050 elements_finished: batch-size vector.
1051 current_input: possibly nested tuple of input tensors.
1052 emit_ta: possibly nested tuple of output TensorArrays.
1053 state: possibly nested tuple of state tensors.
1054 loop_state: possibly nested tuple of loop state tensors.
1055
1056 Returns:
1057 Tuple having the same size as Args but with updated values.
1058 """
1059 (next_output, cell_state) = cell(current_input, state)
1060
1061 nest.assert_same_structure(state, cell_state)
1062 nest.assert_same_structure(cell.output_size, next_output)
1063
1064 next_time = time + 1
1065 (next_finished, next_input, next_state, emit_output,
1066 next_loop_state) = loop_fn(
1067 next_time, next_output, cell_state, loop_state)
1068
1069 nest.assert_same_structure(state, next_state)
1070 nest.assert_same_structure(current_input, next_input)
1071 nest.assert_same_structure(emit_ta, emit_output)
1072
1073 # If loop_fn returns None for next_loop_state, just reuse the
1074 # previous one.
1075 loop_state = loop_state if next_loop_state is None else next_loop_state
1076
1077 def _copy_some_through(current, candidate):
1078 """Copy some tensors through via array_ops.where."""
1079 def copy_fn(cur_i, cand_i):
1080 with ops.colocate_with(cand_i):
1081 return array_ops.where(elements_finished, cur_i, cand_i)
1082 return nest.map_structure(copy_fn, current, candidate)
1083
1084 emit_output = _copy_some_through(zero_emit, emit_output)
1085 next_state = _copy_some_through(state, next_state)
1086
1087 emit_ta = nest.map_structure(
1088 lambda ta, emit: ta.write(time, emit), emit_ta, emit_output)
1089
1090 elements_finished = math_ops.logical_or(elements_finished, next_finished)
1091
1092 return (next_time, elements_finished, next_input,
1093 emit_ta, next_state, loop_state)
1094
1095 returned = control_flow_ops.while_loop(
1096 condition, body, loop_vars=[

Callers

nothing calls this directly

Calls 3

_copy_some_throughFunction · 0.70
loop_fnFunction · 0.50
writeMethod · 0.45

Tested by

no test coverage detected