issues
search
google-deepmind
/
kfac-jax
Second Order Optimization and Curvature Estimation with K-FAC in JAX.
Apache License 2.0
235
stars
18
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
Incorrectly pytree recognition by KFAC optimizer
#273
Uernd
opened
10 hours ago
0
Stackless yashful
#272
copybara-service[bot]
opened
1 day ago
0
Improving polynomial schedule in the examples codebase so that it works as expected when the initial value is *lower* than the final value.
#271
copybara-service[bot]
closed
1 week ago
0
Model parameters being marked as orphan when using KFAC optimizer
#270
Uernd
opened
1 week ago
0
[kfac-jax] Update graph matching test to support the new "algorithm" tuning parameters for dot_general that will be included in the next JAX release.
#269
copybara-service[bot]
closed
1 week ago
0
- Adding support for the "Schedule-free" method to be used as a wrapper for Optax optimizers in the examples codebase.
#268
copybara-service[bot]
closed
2 weeks ago
0
- Passing stats to _post_param_update_processing in examples code.
#267
copybara-service[bot]
closed
3 weeks ago
0
Move optax interface into the main kfac codebase.
#266
copybara-service[bot]
closed
4 weeks ago
0
Pass the state to the update_polyak function
#265
copybara-service[bot]
closed
1 month ago
0
Separated the `optimizers` module in kfac examples into separate modules
#264
copybara-service[bot]
closed
1 month ago
0
Adding the repeated dense graph patterns.
#263
copybara-service[bot]
closed
1 month ago
0
Split `curvature_blocks.py` module into a package.
#262
copybara-service[bot]
closed
1 month ago
0
Split `curvature_estimator.py` module into a package.
#261
copybara-service[bot]
closed
1 month ago
0
Fix progress off by one.
#260
copybara-service[bot]
closed
1 month ago
0
Add an option to specify a different value function for the preconditioner's curvature estimator.
#259
copybara-service[bot]
closed
1 month ago
0
Add TNT blocks to kfac_jax.
#258
copybara-service[bot]
closed
1 month ago
0
Remove the `TwoKroneckerFactored` class and use the `KroneckerFactored` class instead.
#257
copybara-service[bot]
closed
1 month ago
0
ImportError: cannot import name 'psd_inv_cholesky' from 'kfac_jax._src.utils'
#256
eul8
opened
1 month ago
1
Minor non-functional change.
#255
copybara-service[bot]
closed
1 month ago
0
- Adding handling of jitted functions to graph scanner.
#254
copybara-service[bot]
closed
1 month ago
0
Change the default estimation mode of the curvature estimators to `ggn_curvature_prop`
#253
copybara-service[bot]
closed
1 month ago
0
- Improving the logging of the layer registrations and assigned block classes.
#252
copybara-service[bot]
closed
2 months ago
0
Add precon_power option to KFAC optimizer.
#251
copybara-service[bot]
closed
1 month ago
0
- Adding overridable _post_param_update_processing method to experiment class in examples
#250
copybara-service[bot]
closed
2 months ago
0
- Adding feature to BlockDiagonalCurvature to return undamped diagonal.
#249
copybara-service[bot]
closed
2 months ago
0
Minor changes to typing
#248
copybara-service[bot]
closed
2 months ago
0
Replace Bernoulli distributions with Rademachers
#247
copybara-service[bot]
closed
2 months ago
0
Replace Bernoulli distributions with Rademachers
#246
copybara-service[bot]
closed
2 months ago
0
Fixing bug in log_train_stats_with_polyak_avg_every_n_steps of example code.
#245
copybara-service[bot]
closed
2 months ago
0
Compatibility with `pallas` attention
#244
ae-foster
closed
1 week ago
11
- Updating schedule construction code in examples folder so that it properly detects misspelled argument names.
#243
copybara-service[bot]
closed
3 months ago
0
- Fixing a bug in the logging code for Polyak averaging
#242
copybara-service[bot]
closed
3 months ago
0
Minor refactor
#241
copybara-service[bot]
closed
3 months ago
0
* Simplifying the LossTag machinery.
#240
copybara-service[bot]
closed
4 months ago
0
- Expanding Polyak averaging functionality in examples codebase.
#239
copybara-service[bot]
closed
4 months ago
0
Internal change.
#238
copybara-service[bot]
closed
4 months ago
0
Simplifying the LayerTag Primitive machinary.
#237
copybara-service[bot]
closed
4 months ago
0
Internal cleanup
#236
copybara-service[bot]
closed
4 months ago
1
* Removes the loss tags from any tracer transforms. This removes a "undesirable" side affect that when doing fisher/ggn matrix-vector products, they re-introduce the same loss tags in the graph.
#235
copybara-service[bot]
closed
4 months ago
0
Log different block class assignements in the curvature estimator.
#234
copybara-service[bot]
closed
4 months ago
0
Also report the registered losses from the tracer.
#233
copybara-service[bot]
closed
4 months ago
0
* Fixing all Losses to return everything in non-auxiliary data during flattening, to avoid any tracer leaks when the weight is dynamic.
#232
copybara-service[bot]
closed
4 months ago
0
- Label smoothing turned off by default in imagenet experiments. (Was already off in the configs.)
#231
copybara-service[bot]
closed
4 months ago
0
Fixed a type error post google/jax#21261
#230
copybara-service[bot]
closed
4 months ago
0
[KFAC-JAX] Use strings to enumerate kfac jaxpr debug info result paths.
#229
copybara-service[bot]
closed
4 months ago
1
Add more checks around compute quad model
#228
copybara-service[bot]
closed
4 months ago
0
Remove deprecated `jax.tree_map` calls
#227
copybara-service[bot]
closed
4 months ago
0
* Adding damping to the curvature update methods.
#226
copybara-service[bot]
closed
4 months ago
0
Fix xla optimization issue for computation of rho in KFAC.
#225
copybara-service[bot]
closed
4 months ago
1
Can this be used for Laplace approximation?
#224
blackblitz
opened
4 months ago
1
Next