[jax] single decorator grad cache - Githubissues

luyug / GradCache

Run Effective Large Batch Contrastive Learning Beyond GPU/TPU Memory Constraint

Apache License 2.0

327 stars 19 forks source link

[jax] single decorator grad cache #26

Closed luyug closed 6 months ago

luyug commented 6 months ago

Some enhancements to JAX GradCache.

add a single decorator implementation, @grad_cached
improve package structure
improve gitignore