Sean Moriarity

@Ian-GL If you want to tackle this, `Keras` raises when the dropout rate is not on the interval `[0, 1)`, and I think that's a sensible default!

Thanks for this work! I have opened a corollary to this PR in Axon here: https://github.com/elixir-nx/axon/pull/350 I noticed that right now Nx.Random.uniform does not correctly preserve the requested type, you...

Can you share the output of nvidia-smi?

FWIW, we have pre-built XLA packages here: https://github.com/elixir-nx/xla This solution is pretty specific to our projects and so might not be useful for everyone trying to use XLA, but it...

Too much has changed, and this needs to be revisited with considerations from the other library

@fantypants Hey! Just checking up on the status of this PR? Do you need help or another review?

@vishal-h were you able to get this working?

Thinking about this I am considering doing the following: Creating a struct `%Axon.ModelState{parameters: ..., state: ...}` and then adding methods: ```elixir def update_state(%ModelState{}, state) :: model_state def update_parameters(%ModelState{}, parameters) ::...

For more context, this would prevent model state from being passed through `Axon.Updates` as this is unnecessary. So you'd only pass parameters to those and then update the state elsewhere.

Resolved in #553