Sean Moriarity issues

Results 91 issues of


                                            Sean Moriarity

Add initial FP8 support

This PR adds support for an emulated FP8 in Nx's binary backend and adds support for FP8 operations from XLA. The FP8 type in EXLA is just there: `FP8E5M2` type....

Weight tying

Ambiguous inputs should raise

Related to #582 If your model has multiple inputs but you pass a single tensor it will use that for *every* input in the model. It should raise with an...

Remove namespaces

Crash with mixed precision with new Nx

I'm still trying to determine if this is Nx or Axon

Optionally return gradients / gradient state in train step

Merge model state updates into a single `Axon.ModelState` call

Most `Axon.ModelState` functions should be `defn` or `deftransform`

Implement loops as Streams

This takes the regular loop `run` and builds it based off of this `stream` which will return a loop state per epoch. Basically this should be really flexible in that...

Add NIF for loading custom plugins

WIP