Awni Hannun a54f06b16f Fast RMS Norm (#862) ...

* fast rmsnorm

* no rms gpu

* kernel

* fix shared mem

* looped rms and donation in softmax

* Make the squaring in float32 to avoid underflow

* Fix the default StreamOrDevice for rope and rms_norm in fast

* nits

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

2024-03-21 07:20:54 -07:00

..

layers

Fast RMS Norm (#862)

2024-03-21 07:20:54 -07:00

__init__.py

Common neural network initializers nn.initializers (#456)

2024-01-23 06:47:20 -08:00

init.py

updated calls to use loc &scale (#643)

2024-02-08 09:01:59 -08:00

losses.py

Feat: Add weights argument in BCE Loss and tests (#620)

2024-02-07 09:39:52 -08:00

utils.py

Custom VJP and checkpointing (#541)

2024-01-30 16:04:45 -08:00