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

…ckage __init__ for optax.contrib.

WARNING: This can potentially break the loading of a saved optimizer states if the path of the MechanicState class is tied to the saved data (e.g. via pickling). This kind of pickling issue can be fixed at load time by adding an import stub at the old location, or by adding a `sys.modules` entry for the old path.
PiperOrigin-RevId: 570397989

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

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.