google-deepmind / google-deepmind/optax

Code for DoG is duplicated

Open
#1,509 6 comments 1 reaction 1 assignee Claimed by @emilyfertig View on GitHub
type:others
Dominant language
Python
Stars
2.3k
Forks
369
Avg merge
10h 15m
Merged PRs (30d)
7

Description

There seems to be duplicate code for [DoG](https://arxiv.org/abs/2302.12022) in the codebase:

1. optax/_src/transform.py: [`ScaleByDistanceOverGradientsState`](https://github.com/google-deepmind/optax/blob/daecb91b3f0e5de6d3e763c67362e7ac2737bb24/optax/_src/transform.py#L1343), [`scale_by_distance_over_gradients`](https://github.com/google-deepmind/optax/blob/daecb91b3f0e5de6d3e763c67362e7ac2737bb24/optax/_src/transform.py#L1351)

2. optax/contrib/_dog.py: [`DoGState`](https://github.com/google-deepmind/optax/blob/daecb91b3f0e5de6d3e763c67362e7ac2737bb24/optax/contrib/_dog.py#L37), [`scale_by_dog`](https://github.com/google-deepmind/optax/blob/daecb91b3f0e5de6d3e763c67362e7ac2737bb24/optax/contrib/_dog.py#L46)

These two should be merged together.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.