Update oryx version to 0.2.7 and only expose optax_schedule_free wh…

…en it is available.

Fixes #47

PiperOrigin-RevId: 636641008

bayeux/optimize/__init__.py

@@ -15,7 +15,7 @@
 """Imports from submodules."""
 # pylint: disable=g-importing-member
 # pylint: disable=g-import-not-at-top
-import importlib
+import importlib.util
 
 __all__ = []
 
@@ -44,7 +44,9 @@
   # from bayeux._src.optimize.optax import OptimisticGradientDescent  # pylint: disable=line-too-long
   from bayeux._src.optimize.optax import Radam
   from bayeux._src.optimize.optax import Rmsprop
-  from bayeux._src.optimize.optax import ScheduleFree
+  if importlib.util.find_spec("optax.contrib._schedule_free") is not None:
+    from bayeux._src.optimize.optax import ScheduleFree
+    __all__.append("ScheduleFree")
   from bayeux._src.optimize.optax import Sgd
   from bayeux._src.optimize.optax import Sm3
   from bayeux._src.optimize.optax import Yogi
@@ -67,7 +69,6 @@
       # "Dpsgd",
       "Radam",
       "Rmsprop",
-      "ScheduleFree",
       "Sgd",
       "Sm3",
       "Yogi",

pyproject.toml

@@ -16,7 +16,7 @@ keywords = []
 dependencies = [
   "jax>=0.4.6",
   "tensorflow-probability[jax]>=0.19.0",
-  "oryx>=0.2.5",
+  "oryx>=0.2.7",
   "arviz",
   "optax",
   "optimistix",