Angelos Katharopoulos 0de5988f92 Custom VJP and checkpointing (#541) ...

* Implement custom_vjp and checkpointing
* Add a dependency management primitive
* Change the eval order to deep branches first
* Add graph depth tracking to the array

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

..

nn

Custom VJP and checkpointing (#541)

2024-01-30 16:04:45 -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

optimizers.py

Fix SGD implementation (#473)

2024-01-30 15:50:46 -08:00

utils.py

Minor refactor for tree_map and tree_unflatten (#311)

2023-12-28 20:55:10 -08:00