axon
axon copied to clipboard
Nx-powered Neural Networks
Trying to align more with Flax's implementation which does not split keys but rather generates keys from layer meta
The dropout rate is always expected to be in the range of [0, 1). With these changes, now an exception is raised if this happens, instead of letting it crash...
If I understood correctly #362 , the new output of the transform should be `{%Axon.Loop.State{...}, transformed_state}`, so this PR implements said change and updates the tests. This should be the...
WIP: Need to add tests From our original implementation I added a few features, and there's still more I think we need to determine before merging: 1. `preprocess`/`postprocess` - I...
Requires https://github.com/elixir-nx/nx/issues/765
Requires https://github.com/elixir-nx/nx/issues/645
Right now large model compiler stacktraces are very difficult to interpret. It's almost impossible to tell where your issue is because the compiler offers no straightforward information in the stacktrace....