axon
axon copied to clipboard
Published
20 hours ago •
elixir-nx
elixir-nx
Make RNN creation more explicit
Currently we abstract a lot of what goes on within RNNs away to a single layer call with some parameters. Our API is similar to Keras. I am considering making this API more explicit and (hopefully) more flexible. The idea is to introduce the scan layer which applies a scan operation over a sequence:
Axon.input({nil, 128, 32})
|> Axon.lstm_cell(...)
|> Axon.scan()
From here you can control whether the scan is dynamic (run with while) or static (inlined in the graph), the initial state, etc. One issue to consider before going all in on this approach is #99 and how that ties in. Combining these two without introducing an additional cudnn_rnn layer or similar would require some graph optimizations.
