deepmind/kfac-jax issues and pull requests

#165 - Bump numpy from 1.21 to 1.22.0

Pull Request - State: open - Opened by dependabot[bot] about 1 year ago
Labels: dependencies

#164 - Bump ipython from 7.16.1 to 8.10.0

Pull Request - State: open - Opened by dependabot[bot] about 1 year ago
Labels: dependencies

#163 - Bumping the versions of various sphinx packages so the ReadTheDocs don't fail.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#162 - Fixing minor bug that caused an internal JAX error when both jax_debug_nans and jax_disable_jit are enabled.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#161 - [JAX] Migrate XlaBuilder users to emit direct stablehlo MLIR lowerings.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#160 - KFAC Norm Constraint

Issue - State: open - Opened by Baiyu-Su about 1 year ago

#159 - Updated kfac

Pull Request - State: closed - Opened by LeonGerBerlin about 1 year ago - 1 comment

#158 - - Correcting error in docstring for value_and_grad_func argument

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#157 - - Adding copy()'s that were missing from recently added sync methods

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#156 - Fix JAX version check

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#155 - Changing curvature block multiply method (for power 1) to use factored damping when exact_power=False. Previously it just used regular "exact" damping, which was inconsistent with the behavior for other powers.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#154 - Change the test to be compatible with a soon to come Jax release.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#153 - Update to the new Jax convension regrading `preferred_element_type`.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#152 - Remove chex as a dependency for tests.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#151 - Fixing bug that would sometimes cause an exception for networks with scalar-valued parameters.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago

#150 - Add "weighted_norm"-based clipping for arbitrary vectors.

Pull Request - State: open - Opened by copybara-service[bot] about 1 year ago - 1 comment

#149 - Modify to pass `PreconditionState` to `Preconditioner.as_gradient_transform` as an extra arg

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago - 1 comment

#148 - Add the `multiply_curvature` method in `Preconditioner`.

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago - 1 comment

#147 - Add `rel_grad_norm` and `rel_update_norm` stats logging

Pull Request - State: closed - Opened by copybara-service[bot] about 1 year ago - 1 comment

#146 - Fix `Preconditioner` not to pass the out-dated flags (`pmap_axis_name`, `sync`) to `self.estimator.update_curvature_matrix_estimate()`

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#145 - Update RTD setup to python 3.11.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#144 - Fixing a concretization jax error introduced in previous commit.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#143 - * Adding more extensive exponential decay schedule.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#142 - division by zero error for scalar params

Issue - State: open - Opened by szbernat over 1 year ago

#141 - - Adding capability for some schedules to use direct measure of data seen when batch size is unavailable

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#140 - Add options to include grad_, update_, and param_norm to stats.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#139 - Adding automatic cross-device averaging of auxiliary loss/models stats to optimizer.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#138 - - Removing some uses of lru_cache in example experiments that could lead to bugs when certain attributes were changed after construction.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#137 - * Making all datasets to transfer to device on the iterator side before py_prefetch.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#136 - Add Optax-compatible K-FAC `Preconditioner` class

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#135 - * Adding a fix to always fold the process index into the iterator seed.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#134 - * Abstracting away the batch size calculation to a separate class.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#133 - Mod to synchronize matrices in `CurvatureBlock.update_curvature_matrix_estimate()` only when `state.step_counter % inverse_update_period == 0`

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#132 - - Adding error handling for when batch size isn't passed to schedule functions

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#131 - Value functions returning state objects not supported

Issue - State: open - Opened by rmclarke over 1 year ago - 2 comments

#130 - *Minor fixes in printing.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#129 - Add the `_curvature_update_period` flag which controls the timings to perform `_update_estimator_curvature()` in Optimizer.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#128 - Allowing for any kind of optax optimizer to be run if the optimizer is not KFAC.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#127 - Changing the cosine schedule to be a bit more flexible.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#126 - Fix `split` that will be passed to `_build_eval_input()`

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#125 - Adding logging for the number of parameters and optimizer state.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#124 - * Bumping up the package version due to changes in the curvature block classes.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#123 - * Introducing a more general KroneckerFactored block class.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#122 - * Modifying the structure of the experiments class by separating the Supervised experiment from the Jaxline experiment inheritance.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#121 - Making KFAC use `pjit` instead of `pmap`

Pull Request - State: open - Opened by copybara-service[bot] over 1 year ago

#120 - Adding a dense pattern with reshape rather than broadcast on the bias - #114.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#119 - Bumping version to 0.0.4 since we no longer support python 3.7.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#118 - * Updating python annotation to use jax.Array correctly, rather than chex.Array.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#117 - Changing deprecated usages from Jax.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#116 - Reenable some types that were previously disabled due to pytype crashes.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#115 - Reenable some types that were previously disabled due to pytype crashes.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#114 - Quickstart example with different NN libraries does not tag correctly

Issue - State: open - Opened by Binbose over 1 year ago - 6 comments

#113 - - Removing the synchronization across devices of reductions computed in various spots. Reductions ops used to be nondeterministic on the GPU in XLA, but apparently not anymore.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#112 - Suppress pytype errors.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#111 - Fix pytype failures related to teaching pytype about NumPy scalar types.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#110 - Replaces references to jax.numpy.DeviceArray with jax.Array.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#109 - Silence some pytype errors.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#108 - - Minor fixes/improvements to optimizer docstrings

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#107 - * Minor fix when debug is enabled.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#106 - Fixing break in examples code for eval workers with non-K-FAC optimizers.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#105 - Handle checking if a variable is an array-like type for python 3.9 and earlier.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#104 - Fix or ignore some pytype errors related to jnp.ndarray == jax.Array.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#103 - Using kfac inside jitted function

Issue - State: open - Opened by Honza9723 over 1 year ago - 1 comment

#102 - - Modifying examples to only use label smoothing and L2 reg loss when training

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#101 - Adding data seen to the reported statistics on the evaluator in the examples.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#100 - Allowing to register softmax cross entropy without labels.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#99 - Fixing a small bug, where in certain cases the graph matcher can end up in an infinite loop with manual registrations.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#98 - Avoid lhs_shape & rhs_shape parameters in conv_general_dilated

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#97 - Add logging of scaled_grad_norm_sq to track effect of norm constraint

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#96 - fix error when not using multiple devices

Pull Request - State: closed - Opened by y1xiaoc over 1 year ago - 1 comment

#95 - Quick question on "layer_tag_vjp"

Issue - State: closed - Opened by connection-on-fiber-bundles over 1 year ago - 2 comments

#94 - feat(ci): enable `pip` caching in CI

Pull Request - State: closed - Opened by SauravMaheshkar over 1 year ago - 1 comment

#93 - Use jax.device_put_sharded to shard arrays across local devices.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#92 - - Removing unused put_stop_grad_on_loss_factor argument in multiply_fisher_factor

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#91 - OptimizerState is now a type alias for Optimizer.State

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#90 - Fixing bug in MultiChunkAccumulator introduced in recent CL.

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago

#89 - Use jax.device_put_replicated to broadcast to local devices

Pull Request - State: closed - Opened by copybara-service[bot] over 1 year ago - 1 comment

#88 - Adding whitespace to tracer.py to improve readability.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#87 - * Modularizing the utilities file into a separate sub-package.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#86 - Changing examples code so that dataset functions directly return iterators instead of TF datasets.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#85 - Adding inputs_tangent computation to tracer.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#84 - - Improving docstring for optimizer. In particular regarding the damping parameter and LR/momentum/damping adaptation methods.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#83 - Silence some pytype errors.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#82 - Updates the code to always create variables and computations of the same dtype as the its inputs. Previously, if float64 was enabled, some of the results would be (potentially incorrectly) promoted to higher precision.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#81 - - Adding normalization modes feature to classifier loss.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#80 - * Minor fix of how jax scopes are named (':' is not valid).

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#79 - Minor tweaks to training.py (mostly whitespace).

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#78 - * Adding jax.name_scope calls to the important parts of the code for better debugging.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#77 - Minor cleanup of optimizer code and its docstring.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#76 - Fixing minor type error.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#75 - Correct return type annotation for BlockDiagonalCurvature.params_vector_to_blocks_vectors.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago - 1 comment

#74 - * Correct instantiation types for integers in the optimizer state.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#73 - * Minor fixing in handling manual graph registrations.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#72 - Adding white space to some files to improve readability.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#71 - * Minor fix of pi-adjusted-inverse

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#70 - * Fix a bug in the pi-adjusted-inverse.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#69 - * Makng the scale in the curvature blocks to allow for dependance on the state.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#68 - * Making the estimator finalize itself automatically.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#67 - - Using jnp.int64 for data_seen and step counters to avoid overflow

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

#66 - Better PSM tests.

Pull Request - State: closed - Opened by copybara-service[bot] almost 2 years ago

GitHub / deepmind/kfac-jax issues and pull requests