Code for DoG is duplicated

[carlosgmartin]

There seems to be duplicate code for DoG in the codebase:

1. optax/_src/transform.py: ScaleByDistanceOverGradientsState, scale_by_distance_over_gradients

2. optax/contrib/_dog.py: DoGState, scale_by_dog

These two should be merged together.

[abheeyeee]

@carlosgmartin can you review this?

[abheeyeee]

@rajasekharporeddy This PR will fix the issue

[abheeyeee]

Hi @vroulet could you please take a look when you have a moment? Thanks. Also does deepmind have communication channel? slack or discord?

[emilyfertig]

Hi @abheeyeee , I just talked about this with @vroulet . Your PR demonstrates that there isn't much code overlap between the functions -- the code structure and math is the same, but one maps over pytrees and the other reduces, so it isn't possible to share much code (hence all of the branching on layer_wise). Our preference is to keep them as separate functions and move scale_by_distance_over_gradients into contrib and rename it scale_by_l_dog (l for layerwise). Let us know if you're up for making this change -- your refactoring of the scale_by_distance_over_gradients (or layer_wise=True) code is nice, since it highlights the mathematical similarity with the non-layerwise version.

[abheeyeee]

@emilyfertig so what i have understood is that Adding scale_by_l_dog to Contrib, i will create a separate scale_by_l_dog (layerwise DoG) function in the contrib module based on the refactored layer-wise implementation, and update scale_by_distance_over_gradients to be a deprecated wrapper around it.

Proposed Changes
(_dog.py)
Add a new scale_by_l_dog function that implements the layerwise DoG algorithm. This will be a standalone implementation similar to the refactored code, highlighting the mathematical similarity with the non-layerwise version while mapping over pytrees.

[MODIFY] transform.py
Update
scale_by_distance_over_gradients to: Keep the deprecation decorator pointing to scale_by_l_dog instead of
scale_by_dog Call dog.scale_by_l_dog with scale(global_scale) chaining.

[MODIFY] init.py
Add exports for the new function:
scale_by_l_dog
Potentially add an l_dog wrapper function (similar to dog)
[MODIFY]

_dog_test.py
Add test for scale_by_l_dog to ensure:
Correct initialization
Proper layerwise behavior
Compatibility with
scale_by_distance_over_gradients