mlx/python/mlx
Awni Hannun 5798256fcf
Shapeless compilation for some graphs (#687)
* shapeless compilation for some graphs

* update compile benchmark

* default compile a few activations

* buffer donation

* bugfix

* shapeless fix

* update tests to work for cpu and gpu fusion

* test kwargs

* add kwargs to compile

* Recompile when python arguments change

* no compile for tanh

* some constant tests

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
2024-02-19 21:43:54 -08:00
..
nn Shapeless compilation for some graphs (#687) 2024-02-19 21:43:54 -08:00
optimizers Support LR schedulers (#334) 2024-02-15 11:26:20 -08:00
_reprlib_fix.py copyright + ack 2023-11-30 11:12:53 -08:00
extension.py Add isort pre-commit and run (#68) 2023-12-08 11:31:47 -08:00
py.typed Add py.typed to support PEP-561 (type-hinting) for mlx (#588) 2024-01-31 12:05:42 -08:00
utils.py Adds device context manager (#679) 2024-02-14 14:14:58 -08:00