Modernize JAX compatibility and separate core dependencies from example extras

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+build/

README.md

@@ -59,9 +59,13 @@ Install JAX-PI with the following commands:
 ``` 
 git clone https://github.com/PredictiveIntelligenceLab/jaxpi.git
 cd jaxpi
-pip install .
+pip install ".[examples]"
 ```
 
+The `examples` extra installs the logging, configuration, plotting, and
+scientific-computing dependencies used by the scripts under `examples/`. For a
+minimal library installation, use `pip install .` instead.
+
 ## Quickstart
 
 We use [Weights & Biases](https://wandb.ai/site) to log and monitor training metrics. 
@@ -173,4 +177,3 @@ its relative $L^2$ error, and links to the corresponding model [checkpoints](htt
     }
 
 
-

jaxpi/__init__.py

@@ -1,7 +1,21 @@
-from jaxpi import samplers
-from jaxpi import archs
-from jaxpi import models
-from jaxpi import utils
+from importlib import import_module
+
+
+__all__ = ["archs", "models", "samplers", "utils"]
+_SUBMODULES = frozenset(__all__)
+
+
+def __getattr__(name):
+    """Lazily import public submodules while preserving the package API."""
+    if name in _SUBMODULES:
+        module = import_module(f"{__name__}.{name}")
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(set(globals()) | _SUBMODULES)
 
 __version__ = "0.0.1"
 __author__ = "Sifan Wang"

jaxpi/samplers.py

@@ -4,10 +4,8 @@
 import jax.numpy as jnp
 from jax import random, pmap, local_device_count
 
-from torch.utils.data import Dataset
 
-
-class BaseSampler(Dataset):
+class BaseSampler:
     def __init__(self, batch_size, rng_key=random.PRNGKey(1234)):
         self.batch_size = batch_size
         self.key = rng_key

jaxpi/utils.py

@@ -5,9 +5,9 @@
 import jax
 import jax.numpy as jnp
 from jax import jit, grad
-from jax.tree_util import tree_map
 from jax.flatten_util import ravel_pytree
 
+from flax import jax_utils
 from flax.training import checkpoints
 
 
@@ -32,33 +32,28 @@ def ntk_fn(apply_fn, params, *args):
 
 
 def save_checkpoint(state, workdir, keep=5, name=None):
+    workdir = os.path.abspath(workdir)
+
     # Create the workdir if it doesn't exist.
     if not os.path.isdir(workdir):
         os.makedirs(workdir)
 
     # Save the checkpoint.
     if jax.process_index() == 0:
         # Get the first replica's state and save it.
-        state = jax.device_get(tree_map(lambda x: x[0], state))
+        state = jax.device_get(jax_utils.unreplicate(state))
         step = int(state.step)
         checkpoints.save_checkpoint(workdir, state, step=step, keep=keep)
 
 
 def restore_checkpoint(state, workdir, step=None):
-    # check if passed state is in a sharded state
-    # if so, reduce to a single device sharding
-
-    if isinstance(
-        tree_map(lambda x: jnp.array(x).sharding, jax.tree.leaves(state.params))[0],
-        jax.sharding.PmapSharding,
-    ):
-        state = tree_map(lambda x: x[0], state)
-
-    # ensuring that we're in a single device setting
-    assert isinstance(
-        tree_map(lambda x: jnp.array(x).sharding, jax.tree.leaves(state.params))[0],
-        jax.sharding.SingleDeviceSharding,
-    )
+    workdir = os.path.abspath(workdir)
+
+    # Model states are replicated for pmap. Check the scalar step rather than
+    # relying on JAX sharding implementation classes, which change across JAX
+    # versions.
+    if jnp.ndim(state.step) > 0:
+        state = jax_utils.unreplicate(state)
 
     state = checkpoints.restore_checkpoint(workdir, state, step=step)
     return state

setup.py

@@ -1,12 +1,9 @@
-from setuptools import setup, find_packages
-import os
+from pathlib import Path
 
-_CURRENT_DIR = os.path.abspath(os.path.dirname(__file__))
+from setuptools import find_packages, setup
 
-try:
-    README = open(os.path.join(_CURRENT_DIR, "README.md"), encoding="utf-8").read()
-except IOError:
-    README = ""
+
+README = Path(__file__).with_name("README.md").read_text(encoding="utf-8")
 
 setup(
     name="jaxpi",
@@ -16,22 +13,24 @@
     packages=find_packages(),
     python_requires=">=3.8",
     install_requires=[
-        "absl-py",
         "flax",
         "jax",
-        "jaxlib",
-        "matplotlib",
-        "ml_collections",
         "numpy",
         "optax",
-        "scipy",
-        "wandb",
     ],
     extras_require={
+        "examples": [
+            "absl-py",
+            "matplotlib",
+            "ml_collections",
+            "scipy",
+            "tabulate",
+            "wandb",
+        ],
         "testing": ["pytest"],
     },
     license="Apache 2.0",
     description="A library of PINNs models in JAX Flax.",
-    long_description=open(os.path.join(_CURRENT_DIR, "README.md")).read(),
+    long_description=README,
     long_description_content_type="text/markdown",
 )