aspire-inference 0.1.0a6__tar.gz → 0.1.0a8__tar.gz

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/PKG-INFO

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a6
+Version: 0.1.0a8
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
 Project-URL: Homepage, https://github.com/mj-will/aspire
+Project-URL: Documentation, https://aspire.readthedocs.io/
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -38,7 +39,7 @@ Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse
 
-aspire is a framework for reusing existing posterior samples to obtain new results at a reduced code.
+aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
 
 ## Installation
 
@@ -50,3 +51,16 @@ pip install aspire-inference
 
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
 the package can be imported and used as `aspire`.
+
+## Documentation
+
+See the [documentation on ReadTheDocs][docs].
+
+## Citation
+
+If you use `aspire` in your work please cite the [DOI][DOI] and [paper][paper].
+
+
+[docs]: https://aspire.readthedocs.io/
+[DOI]: https://doi.org/10.5281/zenodo.15658747
+[paper]: https://arxiv.org/abs/2511.04218

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/README.md

@@ -1,6 +1,6 @@
 # aspire: Accelerated Sequential Posterior Inference via REuse
 
-aspire is a framework for reusing existing posterior samples to obtain new results at a reduced code.
+aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
 
 ## Installation
 
@@ -12,3 +12,16 @@ pip install aspire-inference
 
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
 the package can be imported and used as `aspire`.
+
+## Documentation
+
+See the [documentation on ReadTheDocs][docs].
+
+## Citation
+
+If you use `aspire` in your work please cite the [DOI][DOI] and [paper][paper].
+
+
+[docs]: https://aspire.readthedocs.io/
+[DOI]: https://doi.org/10.5281/zenodo.15658747
+[paper]: https://arxiv.org/abs/2511.04218

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/aspire_inference.egg-info/PKG-INFO

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a6
+Version: 0.1.0a8
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
 Project-URL: Homepage, https://github.com/mj-will/aspire
+Project-URL: Documentation, https://aspire.readthedocs.io/
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -38,7 +39,7 @@ Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse
 
-aspire is a framework for reusing existing posterior samples to obtain new results at a reduced code.
+aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
 
 ## Installation
 
@@ -50,3 +51,16 @@ pip install aspire-inference
 
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
 the package can be imported and used as `aspire`.
+
+## Documentation
+
+See the [documentation on ReadTheDocs][docs].
+
+## Citation
+
+If you use `aspire` in your work please cite the [DOI][DOI] and [paper][paper].
+
+
+[docs]: https://aspire.readthedocs.io/
+[DOI]: https://doi.org/10.5281/zenodo.15658747
+[paper]: https://arxiv.org/abs/2511.04218

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/aspire_inference.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@
 LICENSE
 README.md
 pyproject.toml
+readthedocs.yml
 .github/workflows/lint.yml
 .github/workflows/publish.yml
 .github/workflows/tests.yml
@@ -11,6 +12,15 @@ aspire_inference.egg-info/SOURCES.txt
 aspire_inference.egg-info/dependency_links.txt
 aspire_inference.egg-info/requires.txt
 aspire_inference.egg-info/top_level.txt
+docs/Makefile
+docs/api.rst
+docs/conf.py
+docs/entry_points.rst
+docs/examples.rst
+docs/index.rst
+docs/installation.rst
+docs/requirements.txt
+docs/user_guide.rst
 examples/basic_example.py
 examples/smc_example.py
 src/aspire/__init__.py
@@ -42,5 +52,6 @@ tests/test_transforms.py
 tests/test_utils.py
 tests/integration_tests/conftest.py
 tests/integration_tests/test_integration.py
+tests/test_flows/test_flows_core.py
 tests/test_flows/test_jax_flows/test_flowjax_flows.py
 tests/test_flows/test_torch_flows/test_zuko_flows.py

aspire_inference-0.1.0a8/docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

aspire_inference-0.1.0a8/docs/api.rst

@@ -0,0 +1,28 @@
+API Reference
+=============
+
+This section documents the main public classes.
+
+Aspire interface
+----------------
+
+.. autoclass:: aspire.Aspire
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Samples utilities
+-----------------
+
+.. automodule:: aspire.samples
+   :members: Samples, SMCSamples
+   :undoc-members:
+   :show-inheritance:
+
+History objects
+---------------
+
+.. automodule:: aspire.history
+   :members: History, FlowHistory, SMCHistory
+   :undoc-members:
+   :show-inheritance:

aspire_inference-0.1.0a8/docs/conf.py

@@ -0,0 +1,57 @@
+"""Sphinx configuration for the aspire documentation."""
+
+from __future__ import annotations
+
+import os
+import sys
+
+# Ensure the package can be imported when building the docs locally
+ROOT = os.path.abspath("..")
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+import aspire  # noqa: E402
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = "aspire"
+copyright = "2025, Michael J. Williams"
+author = "Michael J. Williams"
+version = aspire.__version__
+release = aspire.__version__
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+]
+
+autodoc_typehints = "description"
+autodoc_member_order = "bysource"
+autosummary_generate = True
+napoleon_google_docstring = False
+napoleon_numpy_docstring = True
+napoleon_preprocess_types = True
+
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+# -- Options for HTML output -------------------------------------------------
+
+html_theme = "sphinx_book_theme"
+html_static_path = ["_static"]
+html_title = "aspire"
+html_theme_options = {
+    "path_to_docs": "docs",
+    "repository_url": "https://github.com/mj-will/aspire",
+    "repository_branch": "main",
+    "use_edit_page_button": True,
+    "use_issues_button": True,
+    "use_repository_button": True,
+    "use_download_button": True,
+}

aspire_inference-0.1.0a8/docs/entry_points.rst

@@ -0,0 +1,29 @@
+Entry Points
+============
+
+Aspire uses Python's entry point system to allow for extensibility and
+integration with external libraries. This mechanism enables users to register
+custom components that can be seamlessly integrated into the Aspire framework.
+
+.. _custom_flows:
+
+Custom Flows
+------------
+
+Aspire supports custom flow implementations via the
+``aspire.flows`` entry point group. To register a new flow backend, define an
+entry point in your ``pyproject.toml`` like so:
+
+.. code-block:: toml
+
+    [project.entry-points."aspire.flows"]
+    myflow = "my_module:MyFlowClass"
+
+
+The specified class must inherit from :class:`aspire.flows.base.Flow` (or one
+of the existing flow wrappers), implement the required methods and define the
+``xp`` attribute which specifies the array namespace.
+You can then select your custom flow by setting
+``flow_backend="myflow"`` when initializing Aspire.
+
+For an example see ``GWFlow`` in ``aspire-gw`` (https://github.com/mj-will/aspire-gw).

aspire_inference-0.1.0a8/docs/examples.rst

@@ -0,0 +1,28 @@
+Examples
+========
+
+The repository ships with runnable scripts that demonstrate typical Aspire
+workflows. Execute them from the examples directory after installing the relevant
+extras.
+
+Sequential Monte Carlo (MiniPCN)
+--------------------------------
+
+.. literalinclude:: ../examples/smc_example.py
+   :language: python
+   :linenos:
+   :lines: 1-80
+   :caption: ``examples/smc_example.py`` (excerpt)
+
+Run the full example:
+
+.. code-block:: console
+
+   $ python smc_example.py
+
+The script demonstrates how to:
+
+- Build contrived mixtures of Gaussians for testing,
+- Fit a Neural Spline Flow to biased initial samples,
+- Run adaptive MiniPCN-SMC via :meth:`aspire.Aspire.sample_posterior`,
+- Plot diagnostics (loss curves, beta schedule, corner plots).

aspire_inference-0.1.0a8/docs/index.rst

@@ -0,0 +1,63 @@
+aspire: Accelerated Sequential Posterior Inference via REuse
+============================================================
+
+``aspire`` is a lightweight framework for reusing existing posterior samples
+and normalizing flows to accelerate Bayesian inference. It focuses on
+practical workflows: fit a flow, adaptively run Sequential Monte Carlo (SMC),
+MCMC or importance samplers, and visualise or export the resulting samples.
+
+Key capabilities
+----------------
+
+- Fit flow-based proposals (PyTorch or JAX backends) with automatic handling of
+  bounded and periodic parameters.
+- Run adaptive SMC (MiniPCN or BlackJAX kernels) and importance sampling with
+  detailed diagnostic histories.
+- Inspect results via convenience helpers for evidence estimates, corner plots,
+  and HDF5/JSON export.
+
+Quick start
+-----------
+
+.. code-block:: python
+
+    import numpy as np
+    from aspire import Aspire
+    from aspire.samples import Samples
+
+    def log_likelihood(samples):
+        x = samples.x
+        return -0.5 * np.sum(x**2, axis=-1)
+
+    def log_prior(samples):
+        return -0.5 * np.sum(samples.x**2, axis=-1)
+
+    init = Samples(np.random.normal(size=(2_000, 4)))
+
+    aspire = Aspire(
+        log_likelihood=log_likelihood,
+        log_prior=log_prior,
+        dims=4,
+        parameters=[f"x{i}" for i in range(4)],
+    )
+    aspire.fit(init, n_epochs=20)
+    posterior = aspire.sample_posterior(
+        sampler="smc",
+        n_samples=500,
+        sampler_kwargs=dict(n_steps=100),
+    )
+
+    posterior.plot_corner()
+
+Use the sections below for environment setup, conceptual guidance, runnable
+examples, and the complete API reference.
+
+.. toctree::
+    :maxdepth: 2
+    :caption: Contents
+
+    installation
+    user_guide
+    examples
+    entry_points
+    api

aspire_inference-0.1.0a8/docs/installation.rst

@@ -0,0 +1,73 @@
+Installation
+============
+
+``aspire`` targets Python 3.10+ and relies on ``numpy``, ``matplotlib``,
+``array-api-compat`` and ``h5py`` for core functionality. Optional extras
+provide tighter integration with popular samplers and flow backends.
+
+Basic setup
+-----------
+
+Install the library from PyPI (note the published name):
+
+.. code-block:: console
+
+   $ python -m pip install aspire-inference
+
+The installed distribution exposes the ``aspire`` import namespace.
+
+Optional extras
+---------------
+
+Additional features can be enabled by installing the relevant extras:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 20 70
+
+   * - Extra
+     - Purpose
+   * - ``scipy``
+     - Access to SciPy utilities used by certain transforms.
+   * - ``jax``
+     - JAX + ``flowjax`` backend for training normalizing flows.
+   * - ``torch``
+     - PyTorch + ``zuko`` backend (default) for normalizing flows and flow matching.
+   * - ``minipcn``
+     - Enables the MiniPCN SMC kernel.
+   * - ``emcee``
+     - Enables the ``emcee`` ensemble sampler integration.
+   * - ``blackjax``
+     - Enables the BlackJAX SMC kernel.
+   * - ``test``
+     - Installs ``pytest`` and coverage helpers for local testing.
+
+Install extras via:
+
+.. code-block:: console
+
+   $ python -m pip install "aspire-inference[torch,minipcn]"
+
+From source
+-----------
+
+Clone the repository and install in editable mode:
+
+.. code-block:: console
+
+   $ git clone https://github.com/mj-will/aspire.git
+   $ cd aspire
+    # (optional) create/activate a virtual environment
+   $ python -m pip install -e ".[torch,minipcn]"
+
+After installation, run the unit test suite to confirm everything is wired up:
+
+.. code-block:: console
+
+   $ python -m pytest
+
+Building the docs locally requires ``sphinx`` and (optionally) the
+``sphinx-``. These are installed automatically when you run
+``python -m pip install -r docs/requirements.txt`` if such a file exists, or
+you can install ``sphinx`` manually before invoking ``make html`` inside the
+``docs`` directory.

aspire_inference-0.1.0a8/docs/requirements.txt

@@ -0,0 +1,2 @@
+sphinx>=7.2
+sphinx-book-theme>=1.1

aspire_inference-0.1.0a8/docs/user_guide.rst

@@ -0,0 +1,128 @@
+User Guide
+==========
+
+This guide walks through the main concepts you will use when combining
+existing samples with new inference runs. It focuses on the high-level Python
+API exposed by :class:`aspire.Aspire`.
+
+Workflow overview
+-----------------
+
+1. **Describe your problem** via ``log_likelihood`` and ``log_prior``
+   callables that accept :class:`aspire.samples.Samples` or compatible
+   objects.
+2. **Package initial draws** with :class:`aspire.samples.Samples` (or
+   :class:`aspire.samples.BaseSamples`) to benefit from consistent typing,
+   plotting helpers, and device-aware conversions.
+3. **Fit a proposal** with :meth:`aspire.Aspire.fit` to learn a proposal
+   tailored to the current posterior (normalising flow by default).
+4. **Sample the posterior** using :meth:`aspire.Aspire.sample_posterior`
+   with either importance sampling or an adaptive SMC kernel (MiniPCN,
+   BlackJAX, or custom samplers).
+5. **Inspect, save, and reuse** the resulting
+   :class:`aspire.samples.Samples`, :class:`aspire.history.History`
+   objects, and the fitted flow.
+
+Working with samples
+--------------------
+
+``aspire`` uses dataclasses defined in :mod:`aspire.samples` to keep sample
+arrays, weights, and evidence estimates together. Key features:
+
+* Automatic conversion between array namespaces (NumPy, JAX, PyTorch) via the
+  ``xp`` argument.
+* Convenience exporters (:meth:`aspire.samples.BaseSamples.to_dict`,
+  :meth:`aspire.samples.BaseSamples.save`) for logging or serialisation.
+* Plotting helpers (:meth:`aspire.samples.BaseSamples.plot_corner`) that
+  integrate with ``corner`` while respecting weights.
+
+When constructing your own samples, provide parameter names to enable labelled
+plots and dataframes. Use :meth:`aspire.samples.Samples.from_samples` to
+switch namespaces or merge multiple runs with
+:meth:`aspire.samples.Samples.concatenate`.
+
+Flows and transforms
+--------------------
+
+Aspire can work with any proposal that implements ``sample_and_log_prob`` and
+``log_prob``; normalising flows remain the default. Flows are defined via
+:class:`aspire.flows.base.Flow` and instantiated by
+:meth:`aspire.Aspire.init_flow`. By default Aspire uses the ``zuko``
+implementation of Masked Autoregressive Flows on top of PyTorch. The flow is
+automatically wrapped with :class:`aspire.transforms.FlowTransform` (or a
+composite of bounded / periodic transforms) so you can work with native
+parameter ranges while still optimising in unconstrained space.
+
+You can choose a backend by setting ``flow_backend="flowjax"`` to leverage JAX
+or by providing a fully constructed ``flow`` instance. When ``flow_matching``
+is enabled, Aspire trains a score-based model instead of a classical density
+estimator (requires the `zuko` backend).
+
+External flow implementations can be plugged in via the
+``aspire.flows`` entry point group. See :ref:`custom_flows` for details.
+
+Sampling strategies
+-------------------
+
+The :meth:`aspire.Aspire.sample_posterior` method orchestrates several
+samplers, grouped below by inference style.
+
+Importance sampling
+~~~~~~~~~~~~~~~~~~~
+
+``importance``
+    Draws independent samples from the fitted flow and reweights them using
+    the provided likelihood/prior functions. Perfect for quick sanity checks
+    or sanity bounds on evidence estimates.
+
+Markov chain Monte Carlo
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+``minipcn``
+    Runs the :class:`aspire.samplers.mcmc.MiniPCN` kernel directly (no SMC
+    temperature ladder). Configure ``n_samples`` and pass MCMC kwargs such as
+    ``n_steps`` or ``step_fn`` via ``sampler_kwargs``.
+``emcee``
+    Uses the :class:`aspire.samplers.mcmc.Emcee` ensemble sampler for
+    gradient-free proposals. Provide ``sampler_kwargs`` like ``nwalkers`` or
+    ``n_steps`` to control the chain length.
+
+Sequential Monte Carlo
+~~~~~~~~~~~~~~~~~~~~~~
+
+``smc`` / ``minipcn_smc``
+    Runs adaptive SMC with the MiniPCN MCMC kernel. Configure the number of
+    particles via ``n_samples`` and pass kernel settings in ``sampler_kwargs``
+    (for example ``n_steps``, ``target_acceptance_rate`` or ``step_fn``).
+``blackjax_smc``
+    Uses BlackJAX kernels (requires the ``blackjax`` extra) while keeping the
+    same adaptive temperature schedule as the MiniPCN backend.
+``emcee_smc``
+    Replaces the internal MCMC move with the ``emcee`` ensemble sampler,
+    providing a gradient-free option that still benefits from SMC tempering.
+
+You can plug in custom preconditioning by setting ``preconditioning`` to
+``"standard"`` (affine normalisation based on current samples), ``"flow"``
+(use the fitted flow as a transport map), or ``None`` to disable additional
+transforms.
+
+History, diagnostics, and persistence
+-------------------------------------
+
+Every sampler attaches a history object (see :mod:`aspire.history`) with
+diagnostic metrics such as effective sample size, intermediate temperatures,
+or acceptance rates. Plot them via :meth:`aspire.history.SMCHistory.plot` or
+specialised helpers like :meth:`aspire.history.SMCHistory.plot_beta`.
+
+Use the following methods to persist and later resume work:
+
+* :meth:`aspire.Aspire.save_flow` / :meth:`aspire.Aspire.load_flow` to
+  snapshot the trained flow.
+* :meth:`aspire.Aspire.save_config` or
+  :meth:`aspire.Aspire.save_config_to_json` to capture all hyperparameters.
+* :meth:`aspire.samples.BaseSamples.save` to store weighted samples in HDF5
+  for downstream analysis.
+
+Together these utilities support iterative workflows where you continuously
+refine the proposal distribution, reuse expensive likelihood evaluations, and
+relaunch SMC runs with minimal boilerplate.

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/pyproject.toml

@@ -55,6 +55,7 @@ test = [
 
 [project.urls]
 Homepage = "https://github.com/mj-will/aspire"
+Documentation = "https://aspire.readthedocs.io/"
 
 [tool.setuptools_scm]
 

aspire_inference-0.1.0a8/readthedocs.yml

@@ -0,0 +1,15 @@
+version: 2
+
+sphinx:
+  configuration: docs/conf.py
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.13"
+
+python:
+  install:
+    - method: pip
+      path: .
+    - requirements: docs/requirements.txt

aspire_inference-0.1.0a8/src/aspire/flows/__init__.py

@@ -0,0 +1,68 @@
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def get_flow_wrapper(
+    backend: str = "zuko", flow_matching: bool = False
+) -> tuple[type, Any]:
+    """Get the wrapper for the flow implementation.
+
+    Parameters
+    ----------
+    backend : str
+        The backend to use. Options are "zuko" (PyTorch), "flowjax" (JAX), or
+        any other registered flow class via entry points. Default is "zuko".
+    flow_matching : bool, optional
+        Whether to use flow matching variant of the flow. Default is False.
+
+    Returns
+    -------
+    FlowClass : type
+        The flow class corresponding to the specified backend.
+    xp : module
+        The array API module corresponding to the specified backend.
+    """
+    from importlib.metadata import entry_points
+
+    if backend == "zuko":
+        import array_api_compat.torch as torch_api
+
+        from .torch.flows import ZukoFlow, ZukoFlowMatching
+
+        if flow_matching:
+            return ZukoFlowMatching, torch_api
+        else:
+            return ZukoFlow, torch_api
+    elif backend == "flowjax":
+        import jax.numpy as jnp
+
+        from .jax.flows import FlowJax
+
+        if flow_matching:
+            raise NotImplementedError(
+                "Flow matching not implemented for JAX backend"
+            )
+        return FlowJax, jnp
+    else:
+        if flow_matching:
+            logger.warning(
+                "Flow matching option is ignored for external backends."
+            )
+        eps = {
+            ep.name.lower(): ep for ep in entry_points(group="aspire.flows")
+        }
+        if backend in eps:
+            FlowClass = eps[backend].load()
+            xp = getattr(FlowClass, "xp", None)
+            if xp is None:
+                raise ValueError(
+                    f"Flow class {backend} does not define an `xp` attribute"
+                )
+            return FlowClass, xp
+        else:
+            known_backends = ["zuko", "flowjax"] + list(eps.keys())
+            raise ValueError(
+                f"Unknown backend '{backend}'. Available backends: {known_backends}"
+            )

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/src/aspire/samples.py

@@ -14,6 +14,7 @@ from array_api_compat import (
 )
 from array_api_compat import device as api_device
 from array_api_compat.common._typing import Array
+from matplotlib.figure import Figure
 
 from .utils import (
     asarray,
@@ -161,7 +162,24 @@ class BaseSamples:
 
         return pd.DataFrame(self.to_dict(flat=flat))
 
-    def plot_corner(self, parameters: list[str] | None = None, **kwargs):
+    def plot_corner(
+        self,
+        parameters: list[str] | None = None,
+        fig: Figure | None = None,
+        **kwargs,
+    ):
+        """Plot a corner plot of the samples.
+
+        Parameters
+        ----------
+        parameters : list[str] | None
+            List of parameters to plot. If None, all parameters are plotted.
+        fig : matplotlib.figure.Figure | None
+            Figure to plot on. If None, a new figure is created.
+        **kwargs : dict
+            Additional keyword arguments to pass to corner.corner(). Kwargs
+            are deep-copied before use.
+        """
         import corner
 
         kwargs = copy.deepcopy(kwargs)
@@ -173,7 +191,7 @@ class BaseSamples:
             x = self.x[:, indices] if self.x.ndim > 1 else self.x[indices]
         else:
             x = self.x
-        fig = corner.corner(to_numpy(x), **kwargs)
+        fig = corner.corner(to_numpy(x), fig=fig, **kwargs)
         return fig
 
     def __str__(self):

{aspire_inference-0.1.0a6 → aspire_inference-0.1.0a8}/src/aspire/transforms.py

@@ -684,7 +684,7 @@ class FlowPreconditioningTransform(BaseTransform):
             self.flow_kwargs.setdefault("dtype", dtype)
         self.fit_kwargs = dict(fit_kwargs or {})
 
-        FlowClass = get_flow_wrapper(
+        FlowClass, xp = get_flow_wrapper(
             backend=flow_backend, flow_matching=flow_matching
         )
         transform = CompositeTransform(
@@ -695,7 +695,7 @@ class FlowPreconditioningTransform(BaseTransform):
             bounded_transform=bounded_transform,
             affine_transform=affine_transform,
             device=device,
-            xp=FlowClass.xp,
+            xp=xp,
             eps=eps,
             dtype=dtype,
         )

aspire_inference-0.1.0a8/tests/test_flows/test_flows_core.py

@@ -0,0 +1,47 @@
+import pytest
+
+from aspire.flows import get_flow_wrapper
+
+
+def test_get_flow_wrapper_zuko():
+    FlowClass, xp = get_flow_wrapper(backend="zuko", flow_matching=False)
+    import array_api_compat.torch as torch_api
+
+    from aspire.flows.torch.flows import ZukoFlow
+
+    assert FlowClass is ZukoFlow
+    assert xp is torch_api
+
+
+def test_get_flow_wrapper_zuko_flow_matching():
+    FlowClass, xp = get_flow_wrapper(backend="zuko", flow_matching=True)
+    import array_api_compat.torch as torch_api
+
+    from aspire.flows.torch.flows import ZukoFlowMatching
+
+    assert FlowClass is ZukoFlowMatching
+    assert xp is torch_api
+
+
+def test_get_flow_wrapper_flowjax():
+    FlowClass, xp = get_flow_wrapper(backend="flowjax", flow_matching=False)
+    import jax.numpy as jnp
+
+    from aspire.flows.jax.flows import FlowJax
+
+    assert FlowClass is FlowJax
+    assert xp is jnp
+
+
+def test_get_flow_wrapper_flowjax_flow_matching_not_implemented():
+    with pytest.raises(
+        NotImplementedError,
+        match="Flow matching not implemented for JAX backend",
+    ):
+        get_flow_wrapper(backend="flowjax", flow_matching=True)
+
+
+def test_get_flow_wrapper_external_backend_not_found():
+    msg = "Unknown backend 'unknown_backend'. Available backends: "
+    with pytest.raises(ValueError, match=msg):
+        get_flow_wrapper(backend="unknown_backend")

aspire_inference-0.1.0a6/src/aspire/flows/__init__.py

@@ -1,40 +0,0 @@
-def get_flow_wrapper(backend: str = "zuko", flow_matching: bool = False):
-    """Get the wrapper for the flow implementation."""
-    if backend == "zuko":
-        import array_api_compat.torch as torch_api
-
-        from .torch.flows import ZukoFlow, ZukoFlowMatching
-
-        if flow_matching:
-            return ZukoFlowMatching, torch_api
-        else:
-            return ZukoFlow, torch_api
-    elif backend == "flowjax":
-        import jax.numpy as jnp
-
-        from .jax.flows import FlowJax
-
-        if flow_matching:
-            raise NotImplementedError(
-                "Flow matching not implemented for JAX backend"
-            )
-        return FlowJax, jnp
-    else:
-        from importlib.metadata import entry_points
-
-        eps = {
-            ep.name.lower(): ep
-            for ep in entry_points().get("aspire.flows", [])
-        }
-        if backend in eps:
-            FlowClass = eps[backend].load()
-            xp = getattr(FlowClass, "xp", None)
-            if xp is None:
-                raise ValueError(
-                    f"Flow class {backend} does not define an `xp` attribute"
-                )
-            return FlowClass, xp
-        else:
-            raise ValueError(
-                f"Unknown flow class: {backend}. Available classes: {list(eps.keys())}"
-            )