google-deepmind/kfac-jax
Second Order Optimization and Curvature Estimation with K-FAC in JAX.
PythonApache-2.0
Issues
- 1
Can this be used for Laplace approximation?
#224 opened by blackblitz - 1
- 10
Compatibility with `pallas` attention
#244 opened by ae-foster - 4
- 3
Using K-FAC with physics-based losses
#186 opened by SNMS95 - 4
- 2
KFAC Norm Constraint
#160 opened by Baiyu-Su - 2
Using kfac inside jitted function
#103 opened by Honza9723 - 7
- 1
division by zero error for scalar params
#142 opened by szbernat - 2
Quick question on "layer_tag_vjp"
#95 opened by connection-on-fiber-bundles - 3
- 4
TypeError: 'ShapedArray' object is not iterable
#10 opened by ltz0120 - 4