fix: make scale_by_distance_over_gradients a deprecated wrapper aroun…

[Mahajan-Sachin]

fix: make scale_by_distance_over_gradients a deprecated wrapper around scale_by_l_dog (#1509)
Fixes #1509

Refactor scale_by_distance_over_gradients to delegate to
contrib.scale_by_l_dog, removing duplicate implementation.

Changes:

• Add scale_by_l_dog and l_dog to optax/contrib/_dog.py
• Export from optax/contrib/init.py
• Replace implementation in _src/transform.py with a deprecated wrapper
• Remove duplicate ~70 lines of implementation

Note:
The returned state type changes from
ScaleByDistanceOverGradientsState to LDoGState
(fields are identical).

All tests pass locally (580 passed).

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[rdyro]

Hey, thanks for the contribution attempt!

The formatting/diff tool you're using to make changes appears confused about the indentation width in our files (please use indent-width=2). It'll be much easier to pass the CI if you try to make the diff as small as possible.

You can also run much of the CI locally, by installing pre-commit and running pre-commit run -a in the root of the repository.

[Mahajan-Sachin]

All checks are passing now. Please let me know if any changes are needed. Happy to update further. Thanks!

[rdyro]

We probably want to keep the verbose name "distance_over_gradients" and alias "dog" to it rather than the other way around.

[Mahajan-Sachin]

I’ve applied the suggested changes and kept distance_over_gradients as the canonical name.
All checks are passing now. Please let me know if anything else should be adjusted.

[rdyro]

Please don't import contrib in the core package, retain the implementation from the core package if you want to do the refactor.

[Mahajan-Sachin]

All CI checks passed (lint, pre-commit, pytype, doctests, pytest matrix across versions).
No failures. Ready for review and merge.
Thanks for the feedback earlier – addressed all points (no contrib import in core, kept canonical name, deprecated wrapper retained, layer-wise in contrib).
Happy to make any last tweaks! 🙌

[rdyro]

Your PR appears to copy the implementation from

optax/optax/_src/transform.py

Line 1387 in 1419a82

def scale_by_distance_over_gradients(

into the contrib module.

Also, the conclusion in the issue #1509 is that the code duplication is most likely necessary.

I'm confused by what this PR is trying to achieve.

[Mahajan-Sachin]

Hi @rdyro, thanks again for the feedback!

Just to make sure we're aligned on #1509's original intent:

• @emilyfertig (issue opener and self-assignee in Dec 2025) suggested moving the layer-wise impl to contrib as scale_by_l_dog, and making the core scale_by_distance_over_gradients a deprecated wrapper around it (for backward compat + deprecation warning), while keeping duplication where necessary due to different control flow.

I you can check the thread directly:
#1509

This PR implements exactly that approach:

• New layer-wise code added in contrib/_dog.py (scale_by_l_dog, l_dog, LDoGState)
• Core impl replaced with deprecated wrapper (no contrib import, retains old logic + warns to use contrib version)
• Canonical name preserved, ~70 duplicate lines removed safely

All checks green, backward compatible (identical state fields).

If the project now prefers no wrapper / full duplication without deprecation, or any other direction, just let me know – happy to adjust, revert the wrapper, or close the PR accordingly. No problem at all.

Thanks for your time and guidance! 🙏