BatchJAX is a library that allow JAX vmap to be used over lists and objax.ModuleList.
pip install batchjaxSee batchjax_example.ipynb.
-
Notifications
You must be signed in to change notification settings - Fork 0
Batchjax is a library that allows jax's vmap to be used over list and Objax ModuleLists.
License
defaultobject/batchjax
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
About
Batchjax is a library that allows jax's vmap to be used over list and Objax ModuleLists.