-
Notifications
You must be signed in to change notification settings - Fork 23
Pull requests: google-deepmind/kfac-jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Minor bug fix: changing train_step in examples code to take a mean of the stats instead of taking from the first device. Because the optimizer syncs its own stats (like loss), this didn't matter except for stats returned from the kfac_jax optimizer (or Optax optimizers using OptaxWrapper). However, the Polyak averaged loss wasn't actually synced across devices (as its not part of the optimizer anymore), so "loss_polyak" was being reported only for the first device.
#304
opened Nov 27, 2024 by
copybara-service
bot
Loading…
Adding support for keyword arguments to staged methods.
#302
opened Nov 26, 2024 by
copybara-service
bot
Loading…
Bump numpy from 1.21 to 1.22.0
dependencies
Pull requests that update a dependency file
#165
opened Sep 4, 2023 by
dependabot
bot
Loading…
Bump ipython from 7.16.1 to 8.10.0
dependencies
Pull requests that update a dependency file
#164
opened Sep 4, 2023 by
dependabot
bot
Loading…
ProTip!
Add no:assignee to see everything that’s not assigned.