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

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

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/optax/__init__.py b/optax/__init__.py
@@ -14,8 +14,8 @@
 # ==============================================================================
 """Optax: composable gradient processing and optimization, in JAX."""
 
-from optax import contrib
 from optax import experimental
+from optax._src import contrib
 from optax._src.alias import adabelief
 from optax._src.alias import adafactor
 from optax._src.alias import adagrad

diff --git a/optax/contrib/__init__.py → optax/_src/contrib/__init__.py b/optax/contrib/__init__.py → optax/_src/contrib/__init__.py