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