Fix: #1509 Merge duplicate DoG implementations and add layer-wise support by abheeyeee · Pull Request #1518 · google-deepmind/optax

abheeyeee · 2025-12-02T08:04:35Z

Fix Issue: #1509
This PR merges the duplicate Distance over Gradients (DoG) implementations found in optax/contrib/_dog.py and optax/_src/transform.py into a single, unified implementation in optax/_src/dog.py.
Created optax/_src/dog.py which consolidates DoG and DoWG.
The new scale_by_dognow supports a layer_wise argument.
Re-implemented scale_by_distance_over_gradients in optax/_src/transform.py to use the new scale_by_dog with layer_wise=True.
Deprecated scale_by_distance_over_gradients in favor of scale_by_dog.
Updated optax/contrib/_dog.py to be a compatibility shim importing from optax/_src/dog.py.
Added optax/_src/dog_test.py to verify both global and layer-wise behaviors, as well as legacy compatibility.

vroulet · 2025-12-02T19:03:31Z

Can you keep _dog.py and _dog_test.py files, and keep them in the contrib folder?
This will ease the review process.
Try to keep the pr as concise as possible. (It looks rather good otherwise :) )

abheeyeee · 2025-12-02T19:05:55Z

Can you keep _dog.py and _dog_test.py files, and keep them in the contrib folder? This will ease the review process. Try to keep the pr as concise as possible. (It looks rather good otherwise :) )

Thanks for the Feedback, i will do the changes rn.

abheeyeee · 2025-12-02T19:30:33Z

@vroulet moved the files as you asked. This should fix the issue

emilyfertig

Thanks!

abheeyeee · 2025-12-04T10:07:35Z

@emilyfertig I made those changes that you prefered I added the new scale_by_l_dog function. 137d6a7
commit is not visible, i mistakenly did rebased or reset my branch after making it while trying to fix the merge conflict.
Would really love your feedback. thanks

emilyfertig

Thanks! scale_by_l_dog should be the same as scale_by_dog with layer_wise = True, right? Can we remove the layer_wise arg, and make scale_by_l_dog the same as scale_by_dog with layer_wise=True?

abheeyeee · 2025-12-04T19:03:52Z

Thanks! scale_by_l_dog should be the same as scale_by_dog with layer_wise = True, right? Can we remove the layer_wise arg, and make scale_by_l_dog the same as scale_by_dog with layer_wise=True?

Thanks for your feedback @emilyfertig i did as you asked made scale_by_l_dog the same as scale_by_dog with layer_wise=True.
But i am facing some Checks faiure of Pytest versions on ubuntu jax, and i tried solving it but one passes and another fails. can you help me out here. Otherwise all changes have been made and this should fix the issue

emilyfertig

Thanks! I think the pytest failure is unrelated and should clear up if you rebase.

emilyfertig · 2025-12-04T21:59:35Z

+  return _scale_by_dog(
+      init_step=("heuristic", reps_rel),
+      eps=eps,
+      layer_wise=True,


Sorry, what I meant is to please get rid of the layer_wise arg everywhere, and make separate implementations of scale_by_dog and scale_by_l_dog. Does that make sense?

@emilyfertig Refactored the DoG optimizer implementation in optax/contrib/_dog.py to separate the global and layer-wise variants.
Refactored optax/contrib/_dog.py: Removed the internal _scale_by_dog helper function.
Implemented scale_by_dog (global DoG) and scale_by_l_dog (layer-wise DoG) as distinct, standalone functions.
Removed the layer_wise argument from scale_by_dog to enforce clear separation of concerns.

Updated optax/contrib/_dog_test.py:
Renamed test_dog_layer_wise to test_l_dog_vs_dog to reflect the API changes.
Updated comments to remove outdated references to the layer_wise argument.

Verified that all tests pass with pytest optax/contrib/_dog_test.py.

emilyfertig · 2025-12-05T18:51:29Z

+def scale_by_l_dog(
+    reps_rel: jax.typing.ArrayLike = 1e-6,
+    eps: jax.typing.ArrayLike = 1e-8,
+    param_dtype: Optional[jax.typing.DTypeLike] = None,


Please remove the unused param_dtype arg.

emilyfertig · 2025-12-05T18:54:33Z

+
+  def init_fn(params: base.Params) -> DoGState:
+    params_dtype = optax.tree.dtype(params, "lowest")
+    if param_dtype is not None:


Why is this done here but not in scale_by_dog?

emilyfertig · 2025-12-05T18:56:05Z

-        max_dist=jnp.asarray(r_epsilon, dtype=params_dtype),
-        sum_sq_norm_grads=jnp.asarray(0.0, dtype=params_dtype),
+        init_params=optax.tree.cast(params, params_dtype),
+        max_dist=max_dist,


Leave this inlined?

Please revert this so it's inlined.

emilyfertig · 2025-12-05T19:01:11Z

+    with self.assertRaises(AssertionError):
+        test_utils.assert_trees_all_close(updates_global, updates_layer)
+
+  def test_legacy_compatibility(self):


This test doesn't have much point if the scale_by_distance_over_gradients implementation is changed to call scale_by_l_dog. Can you revert scale_by_distance_over_gradients to its former implementation and still deprecate it?

abheeyeee · 2025-12-05T19:29:50Z

@emilyfertig Understood, this is what i am going to do.
Refactor DoG and Revert Legacy Compatibility:
Proposed Changes

optax/contrib/_dog.py
Remove param_dtype argument from scale_by_l_dog.
Remove param_dtype usage in init_fn inside scale_by_l_dog if it's unused or redundant.
Address consistency between scale_by_l_dog and scale_by_dog regarding init_fn and dtype casting.
Revert scale_by_distance_over_gradients to its legacy implementation (likely a standalone implementation instead of calling scale_by_l_dog) but ensure it relies on scale_by_dog.

optax/contrib/_dog_test.py
Ensure test_legacy_compatibility is meaningful after reverting scale_by_distance_over_gradients.

emilyfertig · 2025-12-05T19:55:07Z

    Ivgi et al, `DoG is SGD's Best Friend: A Parameter-Free Dynamic Step Size
    Schedule <https://arxiv.org/pdf/2302.12022.pdf>`_, 2023
  """
+  reps_rel = 1e-6 if reps_rel is None else reps_rel


This is not the original implementation. You appear to be using an LLM. Please check what it outputs before you request a review.

emilyfertig · 2025-12-05T20:26:43Z

-        max_dist=jnp.asarray(r_epsilon, dtype=params_dtype),
-        sum_sq_norm_grads=jnp.asarray(0.0, dtype=params_dtype),
+        init_params=optax.tree.cast(params, params_dtype),
+        max_dist=max_dist,


Please revert this so it's inlined.

emilyfertig · 2025-12-05T20:27:30Z

+  def init_fn(params: base.Params) -> DoGState:
+    params_dtype = optax.tree.dtype(params, "lowest")
+
+    # r_epsilon is already a tree of scalars


Please remove or clarify comment

abheeyeee marked this pull request as draft December 2, 2025 19:29

abheeyeee marked this pull request as ready for review December 2, 2025 19:29

selamw1 reviewed Dec 3, 2025

View reviewed changes

Comment thread optax/contrib/_dog.py Outdated

emilyfertig reviewed Dec 3, 2025

View reviewed changes

Comment thread optax/_src/transform.py

Comment thread optax/contrib/_dog.py Outdated

Comment thread optax/contrib/_dog_test.py Outdated

abheeyeee requested review from emilyfertig and selamw1 December 3, 2025 05:59

abheeyeee force-pushed the Fix1509 branch from 63b2ab5 to 137d6a7 Compare December 4, 2025 09:19

emilyfertig reviewed Dec 4, 2025

View reviewed changes

abheeyeee requested a review from emilyfertig December 4, 2025 19:03

abheeyeee force-pushed the Fix1509 branch 2 times, most recently from 06259a2 to 0894aeb Compare December 4, 2025 20:03

emilyfertig requested changes Dec 4, 2025

View reviewed changes

abheeyeee added 13 commits December 5, 2025 09:17

fixes Code for DoG is duplicated google-deepmind#1509

efb3d0d

fixes Code for DoG is duplicated google-deepmind#1509

400e6f6

Fix ruff E501 line length issues

a83cfe5

Fix Lint check with pylint

2f1e279

fix build and check types

0a5e381

Move dog optimizer to contrib

4e76d91

Update DoG Optimizer Mixed-Precision Initialization

74ebe74

fixed deprecation warning

67c5d8a

removed comment

c073bfd

resolved merge conflict

ed56f75

fix check types with pytype

e2ecd92

scale_by_l_dog the same as scale_by_dog with layer_wise=True

87e48d9

Rebase

0fd18d9

abheeyeee added 3 commits December 5, 2025 09:24

rebase

c0e8d30

Cross-Platform Compatibility Python ubuntu and jax

8c3a8f3

fix fails

2a07331

abheeyeee force-pushed the Fix1509 branch from 0894aeb to 2a07331 Compare December 5, 2025 03:56

separate implementations of scale_by_dog and scale_by_l_dog

51505f8

abheeyeee requested a review from emilyfertig December 5, 2025 04:23

emilyfertig reviewed Dec 5, 2025

View reviewed changes

Refactor DoG and Revert Legacy Compatibility Walkthrough

30cb988

abheeyeee requested a review from emilyfertig December 5, 2025 19:50

emilyfertig reviewed Dec 5, 2025

View reviewed changes

Restore authentic legacy implementation and fix tests

d963f74

abheeyeee requested a review from emilyfertig December 5, 2025 20:24

emilyfertig requested changes Dec 5, 2025

View reviewed changes

Conversation

abheeyeee commented Dec 2, 2025

Uh oh!

vroulet commented Dec 2, 2025

Uh oh!

abheeyeee commented Dec 2, 2025

Uh oh!

abheeyeee commented Dec 2, 2025

Uh oh!

Uh oh!

emilyfertig left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

abheeyeee commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

emilyfertig left a comment

Choose a reason for hiding this comment

Uh oh!

abheeyeee commented Dec 4, 2025

Uh oh!

emilyfertig left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

abheeyeee commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

abheeyeee commented Dec 4, 2025 •

edited

Loading

abheeyeee commented Dec 5, 2025 •

edited

Loading