Align parameter tagging with Flax conventions

Author: thomaspinder

.gitignore

@@ -153,3 +153,4 @@ node_modules/
 
 docs/api
 docs/_examples
+local_libs/

examples/regression.py

@@ -217,7 +217,7 @@
 # this, we use our defined `posterior` and `likelihood` at our test inputs to obtain
 # the predictive distribution as a `Distrax` multivariate Gaussian upon which `mean`
 # and `stddev` can be used to extract the predictive mean and standard deviatation.
-# 
+#
 # We are only concerned here about the variance between the test points and themselves, so
 # we can just copute the diagonal version of the covariance.  We enforce this by using
 # `return_covariance_type = "diagonal"` in the `predict` call.

gpjax/parameters.py

@@ -77,7 +77,14 @@ def __init__(self, value: T, tag: ParameterTag, **kwargs):
         _check_is_arraylike(value)
 
         super().__init__(value=jnp.asarray(value), **kwargs)
-        self.tag = tag
+
+        # nnx.Variable metadata must be set via set_metadata (direct setattr is disallowed).
+        self.set_metadata(tag=tag)
+
+    @property
+    def tag(self) -> ParameterTag:
+        """Return the parameter's constraint tag."""
+        return self.metadata.get("tag", "real")
 
 
 class NonNegativeReal(Parameter[T]):

pyproject.toml

@@ -140,7 +140,7 @@ help = "Check code formatting and style"
 
 # Testing tasks
 [tool.poe.tasks.test]
-cmd = "pytest . -v -n 8 --beartype-packages='gpjax'"
+cmd = "pytest tests -v -n 8 --beartype-packages='gpjax'"
 help = "Run tests with pytest"
 
 [tool.poe.tasks.coverage]