Awni Hannun be98f4ab6b Reduce a little overhead (#871) ...

* some small overhead improvements

* use result_type in rms_norm

* remove release force

* fix + use non-vector version

* revert compile change

* fix ops

* a little more overhead

* a little more cleanup and overhead

2024-03-22 17:29:36 -07:00

..

layers

Reduce a little overhead (#871)

2024-03-22 17:29:36 -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