Add orbax-checkpoint as dependancy

lululxvi · Sep 25, 2023 · 827f686 · 827f686
1 parent 458d880
commit 827f686
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 2 deletions.
diff --git a/deepxde/model.py b/deepxde/model.py
@@ -4,6 +4,7 @@
 from collections import OrderedDict
 
 import numpy as np
+import orbax.checkpoint as ocp
 
 from . import config
 from . import display
@@ -1012,9 +1013,7 @@ def save(self, save_path, protocol="backend", verbose=0):
                 save_path += ".ckpt"
                 self.net.save_weights(save_path)
             elif backend_name == "jax":
-                # Lazy load Orbax to avoid a hard dependancy when using JAX
                 # TODO: identify a better solution that complies with PEP8 
-                import orbax.checkpoint as ocp
                 from flax.training import orbax_utils 
                 save_path += ".ckpt"
                 checkpoint = {

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -1,5 +1,6 @@
 matplotlib
 numpy
+orbax-checkpoint
 scikit-learn
 scikit-optimize>=0.9.0
 scipy

diff --git a/pyproject.toml b/pyproject.toml
@@ -37,6 +37,7 @@ classifiers = [
 dependencies = [
     "matplotlib",
     "numpy",
+    "orbax-checkpoint",
     "scikit-learn",
     "scikit-optimize>=0.9.0",
     "scipy",

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,6 @@
 matplotlib
 numpy
+orbax-checkpoint
 scikit-learn
 scikit-optimize>=0.9.0
 scipy