Add static field to base, update tests to run on 3.11, fix 3.11 compatibility

.github/workflows/tests.yml

@@ -12,7 +12,7 @@ jobs:
       matrix:
         # Select the Python versions to test against
         os: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
       fail-fast: true
     steps:
       - name: Check out the code

examples/deep_kernels.pct.py

@@ -22,7 +22,7 @@
 from jax.config import config
 from jaxtyping import Array, Float
 from scipy.signal import sawtooth
-from simple_pytree import static_field
+from gpjax.base import static_field
 
 from jaxtyping import install_import_hook
 

examples/kernels.pct.py

@@ -17,7 +17,7 @@
 from jax import jit
 from jax.config import config
 from jaxtyping import Array, Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 
 from jaxtyping import install_import_hook
 

gpjax/base/__init__.py

@@ -20,6 +20,7 @@
     meta_flatten,
     meta_leaves,
     meta_map,
+    static_field,
     save_tree,
 )
 from gpjax.base.param import param_field
@@ -31,6 +32,7 @@
     "meta_map",
     "meta",
     "param_field",
+    "static_field",
     "save_tree",
     "load_tree",
 ]

gpjax/base/module.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 
-__all__ = ["Module", "meta_leaves", "meta_flatten", "meta_map", "meta"]
+__all__ = ["Module", "meta_leaves", "meta_flatten", "meta_map", "meta", "static_field"]
 
 from copy import (
     copy,
@@ -33,6 +33,7 @@
     Tuple,
     TypeVar,
     Union,
+    Mapping,
 )
 import jax
 from jax import lax
@@ -48,13 +49,46 @@
 )
 from simple_pytree import (
     Pytree,
-    static_field,
 )
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 Self = TypeVar("Self")
 
 
+def static_field(
+    default: Any = dataclasses.MISSING,
+    *,
+    default_factory: Any = dataclasses.MISSING,
+    init: bool = True,
+    repr: bool = True,
+    hash: Optional[bool] = None,
+    compare: bool = True,
+    metadata: Optional[Mapping[str, Any]] = None,
+):
+
+    metadata = {} if metadata is None else dict(metadata)
+
+    if "pytree_node" in metadata:
+        raise ValueError("Cannot use metadata with `pytree_node` already set.")
+
+    metadata["pytree_node"] = False
+
+    if default is not dataclasses.MISSING and default_factory is not dataclasses.MISSING:
+        raise ValueError("Cannot specify both default and default_factory.")
+
+    if default is not dataclasses.MISSING:
+        default_factory = lambda: default
+
+    return dataclasses.field(
+    default_factory=default_factory,
+    init=init,
+    repr=repr,
+    hash=hash,
+    compare=compare,
+    metadata=metadata,
+    )
+
+
 class Module(Pytree):
     _pytree__meta: Dict[str, Any] = static_field()
 

gpjax/base/param.py

@@ -53,12 +53,17 @@ def param_field(
     metadata["trainable"] = trainable
     metadata["pytree_node"] = True
 
+    if default is not dataclasses.MISSING and default_factory is not dataclasses.MISSING:
+        raise ValueError("Cannot specify both default and default_factory.")
+
+    if default is not dataclasses.MISSING:
+        default_factory = lambda: default
+
     return dataclasses.field(
-        default=default,
-        default_factory=default_factory,
-        init=init,
-        repr=repr,
-        hash=hash,
-        compare=compare,
-        metadata=metadata,
+    default_factory=default_factory,
+    init=init,
+    repr=repr,
+    hash=hash,
+    compare=compare,
+    metadata=metadata,
     )

gpjax/gps.py

@@ -30,7 +30,7 @@
     Float,
     Num,
 )
-from simple_pytree import static_field
+from gpjax.base import static_field
 
 from gpjax.base import (
     Module,

gpjax/kernels/approximations/rff.py

@@ -2,7 +2,7 @@
 
 from jax.random import PRNGKey
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from gpjax.base import param_field

gpjax/kernels/base.py

@@ -34,7 +34,7 @@
     Float,
     Num,
 )
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax.distributions as tfd
 
 from gpjax.base import (

gpjax/kernels/non_euclidean/graph.py

@@ -22,7 +22,7 @@
     Int,
     Num,
 )
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax as tfp
 
 from gpjax.base import param_field

gpjax/kernels/nonstationary/arccosine.py

@@ -18,7 +18,7 @@
 from beartype.typing import Union
 import jax.numpy as jnp
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from gpjax.base import param_field

gpjax/kernels/nonstationary/polynomial.py

@@ -17,7 +17,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from gpjax.base import param_field

gpjax/kernels/stationary/white.py

@@ -17,7 +17,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax.bijectors as tfb
 import tensorflow_probability.substrates.jax.distributions as tfd
 

gpjax/likelihoods.py

@@ -23,7 +23,7 @@
 import jax.numpy as jnp
 import jax.scipy as jsp
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax as tfp
 
 from gpjax.base import (

gpjax/linops/constant_diagonal_linear_operator.py

@@ -22,7 +22,7 @@
 )
 import jax.numpy as jnp
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 
 from gpjax.linops.diagonal_linear_operator import DiagonalLinearOperator
 from gpjax.linops.linear_operator import LinearOperator

gpjax/mean_functions.py

@@ -28,7 +28,7 @@
     Float,
     Num,
 )
-from simple_pytree import static_field
+from gpjax.base import static_field
 
 from gpjax.base import (
     Module,

gpjax/objectives.py

@@ -6,7 +6,7 @@
 import jax.scipy as jsp
 import jax.tree_util as jtu
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax as tfp
 
 from gpjax.base import Module

gpjax/variational_families.py

@@ -20,7 +20,7 @@
 import jax.numpy as jnp
 import jax.scipy as jsp
 from jaxtyping import Float
-from simple_pytree import static_field
+from gpjax.base import static_field
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from gpjax.base import (

tests/test_base/test_module.py

@@ -32,13 +32,13 @@
 import pytest
 from simple_pytree import (
     Pytree,
-    static_field,
 )
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from gpjax.base.module import (
     Module,
     meta,
+    static_field,
 )
 from gpjax.base.param import param_field
 

tests/test_kernels/test_base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from jax.config import config
 import jax.numpy as jnp
@@ -55,7 +55,7 @@ def test_abstract_kernel():
     # Create a dummy kernel class with __call__ implemented:
     @dataclass
     class DummyKernel(AbstractKernel):
-        test_a: Float[Array, "1"] = jnp.array([1.0])
+        test_a: Float[Array, "1"] = field(default_factory = lambda: jnp.array([1.0]))
         test_b: Float[Array, "1"] = param_field(
             jnp.array([2.0]), bijector=tfb.Softplus()
         )

tests/test_linops/test_linear_operator.py

@@ -21,7 +21,7 @@
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import pytest
-from simple_pytree import static_field
+from gpjax.base import static_field
 
 from gpjax.linops.linear_operator import LinearOperator