mlx/python/mlx/nn
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