google / trax

Trax — Deep Learning with Clear Code and Speed
Apache License 2.0
8.07k stars 813 forks source link

Update LSTM layer by adding more capabilities #1708

Closed sunvod closed 2 years ago

sunvod commented 2 years ago

Add capabilities for LSTM layer: 1) getting a complete state RNN (c, h) from the stack as input instead to use a zero initialization . Use initial_state=True parameter for that. Default False. 2) putting a complete state RNN (c, h) in the stack as output instead to discard it. Use return_state=True parameter for that . Default False.