GW-JAX-Team · thomasckng · Jul 14, 2026 · Jul 14, 2026 · Jul 14, 2026
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -27,7 +27,7 @@ jobs:
         run: uv python install ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          uv sync --python ${{ matrix.python-version }} --group test --group cross-validation --group nested-sampling
+          uv sync --python ${{ matrix.python-version }} --group test --group cross-validation
           if [[ "${{ github.ref }}" == "refs/heads/jim-dev" || "${{ github.base_ref }}" == "jim-dev" ]]; then
             uv pip install git+https://github.com/GW-JAX-Team/flowMC.git@flowMC-dev
             uv pip install git+https://github.com/GW-JAX-Team/ripple.git@ripple-dev

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -23,4 +23,4 @@ repos:
     hooks:
     - id: pyright
       files: src/
-      additional_dependencies: ["astropy>=7.0.1", "beartype>=0.19.0", "evosax>=0.2.0", "flowMC>=0.6.1", "gwpy>=3.0.12", "jax>=0.8.2", "jaxtyping>=0.3.3", pytest, "ripplegw>=0.2.1", typing_extensions, "blackjax@git+https://github.com/blackjax-devs/blackjax.git@main", "pydantic>=2", "anesthetic>=2", "tomli-w>=1.0", "typer>=0.12"]
+      additional_dependencies: ["astropy>=7.0.1", "beartype>=0.19.0", "evosax>=0.2.0", "flowMC>=0.6.1", "gwpy>=3.0.12", "jax>=0.8.2", "jaxtyping>=0.3.3", pytest, "ripplegw>=0.2.1", typing_extensions, "blackjax>=1.6", "pydantic>=2", "anesthetic>=2", "tomli-w>=1.0", "typer>=0.12"]
diff --git a/docs/dev/blackjax_followups.md b/docs/dev/blackjax_followups.md
@@ -1,52 +1,18 @@
-# BlackJAX PyPI follow-ups
+# BlackJAX PyPI migration
 
-## What we're waiting on
+BlackJAX 1.6 includes the nested-sampling APIs Jim requires: the nested slice
+sampler, `blackjax.ns.{base,adaptive,utils}`, and
+`blackjax.ns.utils.finalise`.
 
-A new PyPI release of [blackjax-devs/blackjax](https://github.com/blackjax-devs/blackjax) that includes:
+## Completed
 
-- Nested slice sampler (`blackjax.nss`, `blackjax.ns.{base,adaptive,utils}`).
-- `blackjax.ns.utils.finalise` (used by both NS-AW and NSS at run time).
+- Jim requires `blackjax>=1.6` from PyPI and the lockfile resolves that release.
+- The temporary Git source and `nested-sampling` dependency group were removed.
+- NS-AW and NSS import their BlackJAX modules directly; the version-specific
+  import guards were removed.
+- Installation docs, CI, and pre-commit now use the standard dependency.
 
-These modules are already merged into `blackjax-devs/blackjax` main.
-Jim's `[tool.uv.sources]` currently pins to that branch directly.
+## Outstanding external cleanup
 
-## What to undo when a PyPI release ships
-
-### Packaging
-
-- Remove the `[tool.uv.sources]` block in `pyproject.toml` pointing `blackjax = { git = "https://github.com/blackjax-devs/blackjax.git" }`.
-- Drop the `[dependency-groups] nested-sampling` PEP 735 group entirely.
-- Bump the `blackjax>=1.6` pin to whatever release first contains all NS features.
-- Remove the BlackJAX fork owned by GW JAX Team.
-
-### Inline imports → module top
-
-**`src/jimgw/samplers/blackjax/ns_aw.py`**
-
-| Line | Import                                   |
-|------|------------------------------------------|
-| 144  | `from blackjax.ns.utils import finalise` |
-
-**`src/jimgw/samplers/blackjax/nss.py`**
-
-| Line | Import                                   |
-|------|------------------------------------------|
-| 107  | `from blackjax.ns.utils import finalise` |
-
-### Docs
-
-- `docs/installation.md`:
-  - Drop the "BlackJAX nested samplers" section entirely.
-  - `pip install jimgw` is the canonical install instruction.
-- `docs/guides/samplers.md`:
-  - Drop the "BlackJAX samplers" install preamble section.
-  - Restore `pip install jimgw` as the canonical install instruction.
-- `README.md`: any BlackJAX install caveats can be deleted.
-
-### CI
-
-- `.github/workflows/CI.yml`: remove `--group nested-sampling` from the test job. The `[dependency-groups] nested-sampling` block in `pyproject.toml` can be deleted at the same time.
-
-### Tests
-
-- `tests/unit/samplers/test_blackjax_*.py`: the `pytest.importorskip("blackjax")` lines can stay as defense-in-depth but are no longer load-bearing.
+- Delete the temporary `GW-JAX-Team/blackjax` fork once an organization admin
+  explicitly authorizes that irreversible action.
diff --git a/docs/guides/cli.md b/docs/guides/cli.md
@@ -295,12 +295,12 @@ The `type` field selects the backend. Each backend has its own set of tuning par
 
 ### Sampler comparison
 
-| `type` | Algorithm | Evidence | Extra install |
-| --- | --- | --- | --- |
-| `flowmc` | Normalizing-flow MCMC | No | No |
-| `blackjax-smc` | Sequential Monte Carlo | Yes | No |
-| `blackjax-nss` | Nested slice sampling | Yes | `uv sync --group nested-sampling` |
-| `blackjax-ns-aw` | Nested sampling (acceptance-walk) | Yes | `uv sync --group nested-sampling` |
+| `type` | Algorithm | Evidence |
+| --- | --- | --- |
+| `flowmc` | Normalizing-flow MCMC | No |
+| `blackjax-smc` | Sequential Monte Carlo | Yes |
+| `blackjax-nss` | Nested slice sampling | Yes |
+| `blackjax-ns-aw` | Nested sampling (acceptance-walk) | Yes |
 
 ### `type = "flowmc"`
 

diff --git a/docs/guides/samplers.md b/docs/guides/samplers.md
@@ -11,12 +11,12 @@ samples = jim.get_samples()  # dict[str, np.ndarray] keyed by parameter name
 
 ## Sampler overview
 
-| Sampler | Algorithm | Evidence | Extra install | Prior constraint |
-| --- | --- | --- | --- | --- |
-| [flowMC](#flowmc) | normalizing-flow-enhanced MCMC | No | No | None |
-| [NS-AW](#blackjax-ns-aw) | Nested sampling (bilby/dynesty-style acceptance-walk) | Yes | Yes (`nested-sampling`) | Uniform prior; unit-cube sampling space |
-| [NSS](#blackjax-nss) | Nested slice sampling | Yes | Yes (`nested-sampling`) | Normalised prior |
-| [SMC](#blackjax-smc) | Sequential Monte Carlo | Yes | No | Normalised prior |
+| Sampler | Algorithm | Evidence | Prior constraint |
+| --- | --- | --- | --- |
+| [flowMC](#flowmc) | normalizing-flow-enhanced MCMC | No | None |
+| [NS-AW](#blackjax-ns-aw) | Nested sampling (bilby/dynesty-style acceptance-walk) | Yes | Uniform prior; unit-cube sampling space |
+| [NSS](#blackjax-nss) | Nested slice sampling | Yes | Normalised prior |
+| [SMC](#blackjax-smc) | Sequential Monte Carlo | Yes | Normalised prior |
 
 ---
 
@@ -92,21 +92,6 @@ Key parameters:
 
 ---
 
-## BlackJAX nested samplers
-
-The two BlackJAX nested-sampling backends require additional dependencies.
-The required modules are merged into upstream BlackJAX but not yet released on PyPI; install from main with:
-
-```bash
-uv sync --group nested-sampling
-```
-
-This pulls in:
-
-- **blackjax** — pinned to `blackjax-devs/blackjax` until a PyPI release ships the nested-sampling modules.
-
----
-
 ### BlackJAX NS-AW
 
 Nested sampling with a bilby/dynesty-style adaptive differential-evolution acceptance-walk inner kernel.

diff --git a/docs/installation.md b/docs/installation.md
@@ -24,28 +24,3 @@ pip install -e .
 ```
 
 We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python environment. After cloning the repository, run `uv sync` to create a virtual environment with all dependencies installed.
-
-## BlackJAX nested samplers
-
-Jim's BlackJAX nested-sampling backends depend on modules not yet released on PyPI.
-The features are merged into the upstream [blackjax-devs/blackjax](https://github.com/blackjax-devs/blackjax) main branch but await a versioned release.
-
-### With uv (recommended)
-
-Clone the repository if you haven't already, then sync the nested-sampling group:
-
-```bash
-git clone https://github.com/GW-JAX-Team/Jim.git
-cd Jim
-uv sync --group nested-sampling
-```
-
-This installs BlackJAX from `blackjax-devs/blackjax` as declared in `[tool.uv.sources]`.
-
-### From source (pip)
-
-```bash
-pip install "git+https://github.com/blackjax-devs/blackjax.git"
-```
-
-> **Note:** `pip install jimgw[nested-sampling]` will not work until the required features land in a PyPI release. Use one of the commands above instead.
diff --git a/docs/tutorials/gw150914_cli.md b/docs/tutorials/gw150914_cli.md
@@ -206,7 +206,7 @@ jim-run output/GW150914/config.final.toml
 
 ## 5. Next steps
 
-- **Switch to a nested sampler** — change `[sampler]` to `type = "blackjax-ns-aw"` or `type = "blackjax-nss"` (requires `uv sync --group nested-sampling`) to obtain a log-evidence estimate.
+- **Switch to a nested sampler** — change `[sampler]` to `type = "blackjax-ns-aw"` or `type = "blackjax-nss"` to obtain a log-evidence estimate.
 - **Try an injection** — change `[data]` to `type = "injection"` with `injection_parameters` to validate recovery on a known signal. See the [CLI Config Reference](../guides/cli.md) for the full injection syntax.
 - **Enable heterodyning** — add a `[likelihood.heterodyne]` block for a 10–100× likelihood speedup. See the [CLI Config Reference](../guides/cli.md).
 - **Full config reference** — [CLI Config Reference](../guides/cli.md) documents every available field and default.
diff --git a/pyproject.toml b/pyproject.toml
@@ -24,7 +24,7 @@ dependencies = [
     "anesthetic>=2",
     "astropy>=7.0.1",
     "beartype>=0.19.0",
-    "blackjax>=1.4",
+    "blackjax>=1.6",
     "corner>=2.2.3",
     "equinox>=0.11",
     "evosax>=0.2.0",
@@ -70,10 +70,6 @@ doc = [
     "pymdown-extensions",
     "zensical",
 ]
-nested-sampling = [
-    "blackjax",
-]
-
 [project.urls]
 Homepage = "https://github.com/GW-JAX-Team/Jim"
 Documentation = "https://gw-jax-team.github.io/Jim/"
@@ -83,11 +79,6 @@ Changelog = "https://github.com/GW-JAX-Team/Jim/releases"
 
 [tool.uv.sources]
 jimgw = { workspace = true }
-# Pinned to upstream main until a PyPI release ships with the NS submodules:
-#   - blackjax.ns.{base,adaptive,utils}                                     (NS-AW, NSS)
-#   - blackjax.nss                                                          (NSS)
-# When a release is cut, drop this pin and bump blackjax>=<version>.
-blackjax = { git = "https://github.com/blackjax-devs/blackjax.git" }
 
 [build-system]
 requires = ["hatchling", "hatch-vcs"]

diff --git a/src/jimgw/samplers/__init__.py b/src/jimgw/samplers/__init__.py
@@ -6,10 +6,8 @@
 * `SamplerConfig` — discriminated-union annotation of concrete configs.
 * [`build_sampler`][jimgw.samplers.build_sampler] — factory that dispatches to the right concrete class.
 
-The registry uses lazy loaders so that ``import jimgw.samplers`` does not fail
-when an optional backend (e.g. BlackJAX NSS) is not installed; ImportError is
-raised only when the caller actually asks for that backend via
-`build_sampler`.
+The registry uses loaders to defer concrete sampler construction until the
+caller requests a backend via `build_sampler`.
 """
 
 from collections.abc import Callable
@@ -47,8 +45,7 @@ def register_sampler(type_str: str, lazy_loader: Callable[[], SamplerBuilder]) -
     """Register a concrete [`Sampler`][jimgw.samplers.base.Sampler] class under ``type_str``.
 
     ``lazy_loader`` is called (with no args) only when `build_sampler`
-    dispatches to this type — this is how we defer BlackJAX imports until
-    someone actually asks for a BlackJAX sampler.
+    dispatches to this type.
     """
     if type_str in _REGISTRY:
         raise ValueError(
@@ -81,8 +78,6 @@ def build_sampler(
 
     Raises:
         KeyError: If no sampler is registered for ``config.type``.
-        ImportError: If the lazy loader for that type fails (e.g. BlackJAX
-            missing when requesting a BlackJAX sampler).
     """
     type_str = config.type
     if type_str not in _REGISTRY:
@@ -109,29 +104,10 @@ def build_sampler(
 
 register_sampler("blackjax-smc", lambda: BlackJAXSMCSampler)
 
-# --- BlackJAX NS-AW / NSS (require `uv sync --group nested-sampling`) ---
+from jimgw.samplers.blackjax.ns_aw import BlackJAXNSAWSampler  # noqa: E402
 
+register_sampler("blackjax-ns-aw", lambda: BlackJAXNSAWSampler)
 
-def _load_ns_aw() -> SamplerBuilder:
-    import blackjax
-    from jimgw.samplers.blackjax._imports import require_nested_sampling
+from jimgw.samplers.blackjax.nss import BlackJAXNSSSampler  # noqa: E402
 
-    require_nested_sampling(blackjax)
-    from jimgw.samplers.blackjax.ns_aw import BlackJAXNSAWSampler
-
-    return BlackJAXNSAWSampler
-
-
-def _load_nss() -> SamplerBuilder:
-    import blackjax
-    from jimgw.samplers.blackjax._imports import require_nested_sampling, require_nss
-
-    require_nested_sampling(blackjax)
-    require_nss(blackjax)
-    from jimgw.samplers.blackjax.nss import BlackJAXNSSSampler
-
-    return BlackJAXNSSSampler
-
-
-register_sampler("blackjax-ns-aw", _load_ns_aw)
-register_sampler("blackjax-nss", _load_nss)
+register_sampler("blackjax-nss", lambda: BlackJAXNSSSampler)
diff --git a/src/jimgw/samplers/blackjax/__init__.py b/src/jimgw/samplers/blackjax/__init__.py
@@ -1,5 +1 @@
-"""BlackJAX-backed samplers for Jim.
-
-SMC uses ``blackjax>=1.4`` (a core dependency).
-NS-AW and NSS additionally need: ``uv sync --group nested-sampling``
-"""
+"""BlackJAX-backed samplers for Jim."""
diff --git a/src/jimgw/samplers/blackjax/_imports.py b/src/jimgw/samplers/blackjax/_imports.py
diff --git a/src/jimgw/samplers/blackjax/ns_aw.py b/src/jimgw/samplers/blackjax/ns_aw.py
@@ -11,21 +11,17 @@
 import numpy as np
 from jaxtyping import Array, Float, Key
 from anesthetic.samples import NestedSamples
-import blackjax
 from blackjax.ns.adaptive import AdaptiveNSState
 from blackjax.ns.base import NSInfo
 from blackjax.ns.utils import finalise
 
 from jimgw.samplers.base import Sampler
 from jimgw.samplers.blackjax._acceptance_walk_kernel import bilby_adaptive_de_sampler
-from jimgw.samplers.blackjax._imports import require_nested_sampling
 from jimgw.samplers.config import BlackJAXNSAWConfig
 from jimgw.samplers.periodic import to_unit_cube_stepper
 
 logger = logging.getLogger(__name__)
 
-require_nested_sampling(blackjax)
-
 
 class BlackJAXNSAWSampler(Sampler):
     """BlackJAX nested sampler using the bilby/dynesty-style adaptive DE acceptance-walk kernel.

diff --git a/src/jimgw/samplers/blackjax/nss.py b/src/jimgw/samplers/blackjax/nss.py
@@ -19,18 +19,11 @@
 from blackjax.ns.nss import live_covariance, sample_direction_from_covariance
 from blackjax.ns.utils import finalise
 from jimgw.samplers.base import Sampler
-from jimgw.samplers.blackjax._imports import (
-    require_nested_sampling,
-    require_nss,
-)
 from jimgw.samplers.config import BlackJAXNSSConfig
 from jimgw.samplers.periodic import to_prior_space_proposal
 
 logger = logging.getLogger(__name__)
 
-require_nested_sampling(blackjax)
-require_nss(blackjax)
-
 
 class BlackJAXNSSSampler(Sampler):
     """BlackJAX Nested Slice Sampler (NSS).