Add orbax-checkpoint as dependency

Author: tttc3

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 = {

docs/requirements.txt

@@ -1,5 +1,6 @@
 matplotlib
 numpy
+orbax-checkpoint
 scikit-learn
 scikit-optimize>=0.9.0
 scipy

pyproject.toml

@@ -37,6 +37,7 @@ classifiers = [
 dependencies = [
     "matplotlib",
     "numpy",
+    "orbax-checkpoint",
     "scikit-learn",
     "scikit-optimize>=0.9.0",
     "scipy",

requirements.txt

@@ -1,5 +1,6 @@
 matplotlib
 numpy
+orbax-checkpoint
 scikit-learn
 scikit-optimize>=0.9.0
 scipy