Add better parameter filtering (#530)

* Add better parameter filtering

* Clean up conditional

* Bump version

Author: thomaspinder

docs/scripts/gen_examples.py

@@ -15,7 +15,11 @@
 
 
 def process_file(file: Path, out_file: Path | None = None, execute: bool = False):
-    """Converts a python file to markdown using jupytext and nbconvert."""
+    """Converts a python file to markdown using jupytext and nbconvert.
+    
+    Raises:
+        subprocess.CalledProcessError: If the conversion fails.
+    """
 
     out_dir = out_file.parent
     command = f"cd {out_dir.as_posix()} && "
@@ -30,7 +34,11 @@ def process_file(file: Path, out_file: Path | None = None, execute: bool = False
     else:
         command += f"jupytext --to markdown {file} --output {out_file}"
 
-    subprocess.run(command, shell=True, check=False)
+    result = subprocess.run(command, shell=True, check=False, capture_output=True, text=True)
+    if result.returncode != 0:
+        error_msg = f"Failed to process {file.name}: {result.stderr}"
+        print(error_msg)
+        raise subprocess.CalledProcessError(result.returncode, command, output=result.stdout, stderr=result.stderr)
 
 
 def is_modified(file: Path, out_file: Path):
@@ -63,30 +71,51 @@ def main(args):
 
     print(files)
 
+    # Track failures
+    failures = []
+
     # process files in parallel
     if args.parallel:
         with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
-            futures = []
+            futures = {}
             for file in files:
                 out_file = out_dir / f"{file.stem}.md"
-                futures.append(
-                    executor.submit(
-                        process_file, file, out_file=out_file, execute=args.execute
-                    )
+                future = executor.submit(
+                    process_file, file, out_file=out_file, execute=args.execute
                 )
+                futures[future] = file
 
             for future in as_completed(futures):
+                file = futures[future]
                 try:
                     future.result()
+                    print(f"Successfully processed: {file.name}")
                 except Exception as e:
-                    print(f"Error processing file: {e}")
+                    print(f"Error processing {file.name}: {e}")
+                    failures.append((file, e))
     else:
         for file in files:
             out_file = out_dir / f"{file.stem}.md"
-            process_file(file, out_file=out_file, execute=args.execute)
+            try:
+                process_file(file, out_file=out_file, execute=args.execute)
+                print(f"Successfully processed: {file.name}")
+            except Exception as e:
+                print(f"Error processing {file.name}: {e}")
+                failures.append((file, e))
+    
+    # Report failures and exit with error code if any failed
+    if failures:
+        print(f"\n{len(failures)} file(s) failed to process:")
+        for file, error in failures:
+            print(f"  - {file.name}")
+        return 1  # Return non-zero exit code
+    else:
+        print(f"\nAll {len(files)} file(s) processed successfully!")
+        return 0
 
 
 if __name__ == "__main__":
+    import sys
     project_root = Path(__file__).parents[2]
 
     parser = ArgumentParser()
@@ -99,4 +128,5 @@ def main(args):
     parser.add_argument("--parallel", type=bool, default=False)
     args = parser.parse_args()
 
-    main(args)
+    exit_code = main(args)
+    sys.exit(exit_code)

examples/backend.py

@@ -1,3 +1,19 @@
+# ---
+# jupyter:
+#   jupytext:
+#     cell_metadata_filter: -all
+#     custom_cell_magics: kql
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.11.2
+#   kernelspec:
+#     display_name: .venv
+#     language: python
+#     name: python3
+# ---
+
 # %% [markdown]
 # # Backend Module Design
 #
@@ -116,7 +132,7 @@
 # the parameter's value using a tree map operation.
 
 # %%
-print(constant_param._tag)
+print(constant_param.tag)
 
 # %% [markdown]
 # For most users, you will not need to worry about this as we provide a set of default
@@ -126,7 +142,7 @@
 # see how you can define your own bijectors and parameter types.
 
 # %%
-print(DEFAULT_BIJECTION[constant_param._tag])
+print(DEFAULT_BIJECTION[constant_param.tag])
 
 # %% [markdown]
 # We see here that the Softplus bijector is specified as the default for strictly
@@ -229,7 +245,6 @@
 # altering the way in which we invoke `nnx.split`.
 
 # %%
-
 graphdef, positive_reals, other_params = nnx.split(posterior, PositiveReal, ...)
 print(positive_reals)
 

examples/barycentres.py

@@ -50,6 +50,7 @@
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
+    from gpjax.parameters import Parameter
 
 
 key = jr.key(123)
@@ -179,6 +180,7 @@ def fit_gp(x: jax.Array, y: jax.Array) -> npd.MultivariateNormal:
         model=posterior,
         objective=nmll,
         train_data=D,
+        trainable=Parameter,
     )
     latent_dist = opt_posterior.predict(xtest, train_data=D)
     return opt_posterior.likelihood(latent_dist)

examples/classification.py

@@ -40,13 +40,18 @@
 import optax as ox
 
 from examples.utils import use_mpl_style
-from gpjax.linalg import lower_cholesky, PSD, solve
+from gpjax.linalg import (
+    PSD,
+    lower_cholesky,
+    solve,
+)
 
 config.update("jax_enable_x64", True)
 
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
+    from gpjax.parameters import Parameter
 
 
 identity_matrix = jnp.eye
@@ -119,7 +124,6 @@
 
 # %%
 optimiser = ox.adam(learning_rate=0.01)
-
 opt_posterior, history = gpx.fit(
     model=posterior,
     # we use the negative lpd as we are minimising
@@ -128,6 +132,7 @@
     optim=ox.adamw(learning_rate=0.01),
     num_iters=1000,
     key=key,
+    trainable=Parameter,  # train all parameters (default behavior)
 )
 
 # %% [markdown]
@@ -224,7 +229,7 @@
 
 # Negative Hessian,  H = -∇²p_tilde(y|f):
 graphdef, params, *static_state = nnx.split(
-    opt_posterior, gpx.parameters.Parameter, ...
+    opt_posterior, Parameter, ...
 )
 
 

examples/collapsed_vi.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.11.2
+#       jupytext_version: 1.17.3
 #   kernelspec:
 #     display_name: .venv
 #     language: python
@@ -45,6 +45,7 @@
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
+    from gpjax.parameters import Parameter
 
 
 # set the default style for plotting
@@ -137,6 +138,7 @@
 # _optimise_ their location such that the evidence lower bound is maximised.
 
 # %%
+# Use the enhanced fit API with trainable parameter filtering
 opt_posterior, history = gpx.fit(
     model=q,
     # we want want to minimize the *negative* ELBO
@@ -145,6 +147,7 @@
     optim=ox.adamw(learning_rate=1e-2),
     num_iters=500,
     key=key,
+    trainable=Parameter,
 )
 
 # %%

examples/constructing_new_kernels.py

@@ -41,14 +41,14 @@
 from gpjax.parameters import (
     DEFAULT_BIJECTION,
     PositiveReal,
-    Static,
 )
 
 config.update("jax_enable_x64", True)
 
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
+    from gpjax.parameters import Parameter
 
 
 # set the default style for plotting
@@ -249,7 +249,7 @@ def log_abs_det_jacobian(self, x, y, intermediates=None):
 
 
 class Polar(gpx.kernels.AbstractKernel):
-    period: Static
+    period: float
     tau: PositiveReal
 
     def __init__(
@@ -260,13 +260,13 @@ def __init__(
         n_dims: int | None = None,
     ):
         super().__init__(active_dims, n_dims, DenseKernelComputation())
-        self.period = Static(jnp.array(period))
+        self.period = jnp.array(period)
         self.tau = PositiveReal(jnp.array(tau), tag="polar")
 
     def __call__(
         self, x: Float[Array, "1 D"], y: Float[Array, "1 D"]
     ) -> Float[Array, "1"]:
-        c = self.period.value / 2.0
+        c = self.period / 2.0
         t = angular_distance(x, y, c)
         K = (1 + self.tau.value * t / c) * jnp.clip(
             1 - t / c, 0, jnp.inf
@@ -315,6 +315,7 @@ def __call__(
     model=circular_posterior,
     objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
+    trainable=Parameter,
 )
 
 # %% [markdown]

examples/deep_kernels.py

@@ -60,6 +60,9 @@
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
     from gpjax.kernels.base import AbstractKernel
+    from gpjax.parameters import (
+        Parameter,
+    )
 
 
 # set the default style for plotting
@@ -219,13 +222,18 @@ def __call__(self, x: jax.Array) -> jax.Array:
     ox.adamw(learning_rate=schedule),
 )
 
+# Train all parameters (default behavior with trainable=Parameter)
+# Alternative options for selective training:
+# - trainable=PositiveReal  # only train positive parameters
+# - trainable=lambda module, path, value: 'kernel' in path  # only kernel params
 opt_posterior, history = gpx.fit(
     model=posterior,
     objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
     optim=optimiser,
     num_iters=800,
     key=key,
+    trainable=Parameter,  # explicitly specify trainable filter (default)
 )
 
 # %% [markdown]

examples/graph_kernels.py

@@ -43,6 +43,7 @@
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
+    from gpjax.parameters import Parameter
 
 
 # set the default style for plotting
@@ -179,6 +180,7 @@
     model=posterior,
     objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
+    trainable=Parameter,
 )
 
 # %% [markdown]

examples/intro_to_kernels.py

@@ -37,14 +37,14 @@
 from sklearn.preprocessing import StandardScaler
 
 from examples.utils import use_mpl_style
-from gpjax.parameters import Static
 from gpjax.typing import Array
 
 config.update("jax_enable_x64", True)
 
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
+    from gpjax.parameters import Parameter
 
 
 key = jr.key(42)
@@ -264,7 +264,7 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:  # noqa: F821
 prior = gpx.gps.Prior(mean_function=mean, kernel=kernel)
 
 likelihood = gpx.likelihoods.Gaussian(
-    num_datapoints=D.n, obs_stddev=Static(jnp.array(1e-3))
+    num_datapoints=D.n, obs_stddev=jnp.array(1e-3)
 )  # Our function is noise-free, so we set the observation noise's standard deviation to a very small value
 
 no_opt_posterior = prior * likelihood
@@ -281,6 +281,7 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:  # noqa: F821
     model=no_opt_posterior,
     objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
+    trainable=Parameter,
 )
 
 
@@ -546,13 +547,17 @@ def loss(posterior, data):
     return -gpx.objectives.conjugate_mll(posterior, data)
 
 
+# Optimize all parameters. Alternative filtering strategies available:
+# - trainable=gpx.PositiveReal: train only positive parameters
+# - custom filters for specific parameter subsets
 opt_posterior, history = gpx.fit(
     model=posterior,
     objective=loss,
     train_data=D,
     optim=ox.adamw(learning_rate=1e-2),
     num_iters=500,
     key=key,
+    trainable=Parameter,  # train all parameters (default)
 )