- Using version guard to fix change that broke backwards compatibility with some older versions of JAX.

google-deepmind / kfac-jax

Second Order Optimization and Curvature Estimation with K-FAC in JAX.

Apache License 2.0

250 stars 23 forks source link

Closed copybara-service[bot] closed 2 weeks ago

copybara-service[bot] commented 2 weeks ago

Using version guard to fix change that broke backwards compatibility with some older versions of JAX.
Adding test of in_pmap.