Avoid one layer of indirection by colocating the source and the subpa…

…ckage __init__ for optax.contrib. PiperOrigin-RevId: 570045729
google-deepmind · Oct 3, 2023 · 48b5810 · 48b5810
1 parent e43f9f5
commit 48b5810
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/optax/contrib/__init__.py b/optax/contrib/__init__.py
@@ -14,5 +14,5 @@
 # ==============================================================================
 """Contributed optimizers in Optax."""
 
-from optax._src.contrib.mechanic import MechanicState
-from optax._src.contrib.mechanic import mechanize
+from optax.contrib.mechanic import MechanicState
+from optax.contrib.mechanic import mechanize
diff --git a/optax/_src/contrib/mechanic.py → optax/contrib/mechanic.py b/optax/_src/contrib/mechanic.py → optax/contrib/mechanic.py
diff --git a/optax/_src/contrib/mechanic_test.py → optax/contrib/mechanic_test.py b/optax/_src/contrib/mechanic_test.py → optax/contrib/mechanic_test.py
@@ -27,7 +27,7 @@
 from optax._src import numerics
 from optax._src import state_utils
 from optax._src import update
-from optax._src.contrib import mechanic
+from optax.contrib import mechanic
 
 
 # TODO(harshm): make LARS and Fromage work with mechanic.