epgx 0.1.dev0__py3-none-any.whl

epgx/sequences.py

@@ -0,0 +1,850 @@
+"""Parametrized sequences: the `@sequence` decorator (layer 3)
+
+A Python function building a layer-2 sequence (a list of operators) becomes a
+`Sequence`: a parametrized sequence, called to give the signal. Its variables
+are declared in the decorator by their shape inside the function (one
+simulation):
+
+    name        a scalar
+    name[]      a vector of any length
+    name[n]     a vector, dimension n shared with other variables
+    name[n, m]  a matrix
+
+Every other argument is static: it fixes the structure of the sequence and is
+part of the cache key of the compiled program (lists and arrays are made
+hashable; each new value compiles again).
+
+Batch dimensions are not declared: they come from the values. The dimensions
+of a value beyond its declared core are its leading batch dimensions. An
+`Axis` names them (`Axis("candidate", alphas)` of shape (10, 400) for
+`alpha[n]`): values with the same name are paired, different names form a
+grid; unnamed batch dimensions are broadcast as NumPy (aligned on the right)
+and named `batch_0`, `batch_1`... The function is mapped over the batch
+dimensions (`jax.vmap`), so it always sees its variables with their core
+shape. Outputs: the ADC axis first, then the batch axes (order of first
+appearance), then the axes of the sample and scanner.
+
+`sample=` and `scanner=` are reserved arguments, as in `simulate`; simulation
+options (method, probe, group, output...) are set in the decorator or per
+call (the call wins).
+
+Example:
+    @sequence(variables=["alpha", "phi"], method=Lattice(nmax=20))
+    def pspgr(alpha, phi, TR=5.5, TE=2.25, ndummy=500):
+        n = jnp.arange(ndummy + 1)
+        phases = phi * n * (n + 1) / 2
+        rlx1, rlx2 = E(TE), E(TR - TE)
+        dummies = [[T(alpha, ph), rlx1, rlx2, S(1)] for ph in phases[:-1]]
+        return dummies + [T(alpha, phases[-1]), rlx1, Adc(phase=phases[-1])]
+
+    points = dict(alpha=Axis("point", alphas), phi=Axis("point", phis))
+    signal = pspgr(**points, sample=Sample(T1=1220.0, T2=t2))  # (1, point, T2)
+"""
+
+import collections
+import functools
+import inspect
+import re
+import warnings
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from . import axes, diff, functions, stats
+from .axes import Axis, is_axis
+from .functions import OUTPUTS, OutputLayout, _to_numpy, _to_xarray, simulate
+from .sample import Sample
+
+# reserved arguments and simulation options
+RESERVED = ("sample", "scanner", "select")
+OPTIONS = (
+    "probe",
+    "group",
+    "output",
+    "method",
+    "init",
+    "dtype",
+    "dims",
+    "alignment",
+    "max_unroll",
+    "checkpoint",
+    "merge",
+    "jit",
+    "disp",
+)
+FIELDS = Sample.FIELDS + ("B1", "B0")
+# options of `trace` (others, e.g. probe or output, do not apply)
+TRACE_OPTIONS = (
+    "method",
+    "init",
+    "dtype",
+    "dims",
+    "alignment",
+    "max_unroll",
+    "jit",
+    "disp",
+)
+
+# warning after this number of values of a static array argument
+RECOMPILATIONS = 3
+
+_DECLARATION = re.compile(
+    r"\s*([A-Za-z_]\w*)\s*(\[\s*(?:[A-Za-z_]\w*\s*(?:,\s*[A-Za-z_]\w*\s*)*)?\])?\s*"
+)
+
+
+def sequence(function=None, *, variables=(), **options):
+    """decorator: turn a function returning a sequence into a `Sequence`
+
+    Args:
+        function: function returning a (nested) list of operators
+        variables: declarations of the traced arguments ("T2", "alpha[n]"...),
+            cf. module docstring; every other argument is static
+        options: default simulation options (cf. `simulate`)
+
+    Returns:
+        Sequence
+    """
+    if function is None:
+        return functools.partial(sequence, variables=variables, **options)
+    return Sequence(function, variables, **options)
+
+
+def parse_variable(declaration):
+    """name and core dimensions of a variable declaration
+
+    Returns:
+        name, dims: tuple of core dimension names ("name[]" for `name[]`)
+    """
+    match = _DECLARATION.fullmatch(declaration)
+    if not match:
+        raise ValueError(
+            f"Invalid variable declaration: {declaration!r} "
+            "(expected name, name[], name[n] or name[n, m])"
+        )
+    name, brackets = match.groups()
+    if brackets is None:
+        return name, ()
+    dims = tuple(dim.strip() for dim in brackets[1:-1].split(",") if dim.strip())
+    return name, dims or (f"{name}[]",)
+
+
+class Sequence:
+    """Parametrized sequence (cf. `sequence`)
+
+    Attributes:
+        function: function returning the layer-2 sequence
+        variables: dict name -> core dimensions of the declared variables
+        options: default simulation options
+    """
+
+    def __init__(self, function, variables=(), **options):
+        functools.update_wrapper(self, function)
+        if isinstance(variables, str):
+            variables = [variables]
+        self.function = function
+        self.signature = inspect.signature(function)
+        self.variables = dict(parse_variable(v) for v in variables)
+        if len(self.variables) != len(variables):
+            raise ValueError(f"Variables declared twice: {list(variables)}")
+
+        params = self.signature.parameters
+        for name in params:
+            if name in RESERVED + OPTIONS:
+                raise ValueError(
+                    f"Argument {name!r} of {self.__name__} is reserved by @sequence"
+                )
+        for name in self.variables:
+            if name not in params:
+                raise ValueError(
+                    f"Unknown variable {name!r}: not an argument of "
+                    f"{self.__name__}{self.signature}"
+                )
+            elif name in FIELDS:
+                raise ValueError(
+                    f"Variable {name!r} reuses the name of a sample/scanner field"
+                )
+        unknown = set(options) - set(OPTIONS)
+        if unknown:
+            raise ValueError(f"Unknown option(s): {sorted(unknown)}")
+        self.options = options
+        self._seen = {}  # static array argument -> values seen (recompilations)
+        self._compiled = jax.jit(
+            self._execute, static_argnames=("statics", "options", "kinds")
+        )
+
+    def __repr__(self):
+        variables = [_declaration(name, dims) for name, dims in self.variables.items()]
+        return f"Sequence({self.__name__}, variables={variables})"
+
+    def __call__(self, *args, sample=None, scanner=None, **kwargs):
+        """signal (cf. `signal`)"""
+        return self.signal(*args, sample=sample, scanner=scanner, **kwargs)
+
+    def signal(self, *args, sample=None, scanner=None, **kwargs):
+        """simulated signal
+
+        Args:
+            args, kwargs: arguments of the function (variables: with optional
+                batch dimensions) and simulation options (cf. `simulate`)
+            sample, scanner: cf. `simulate`
+
+        Returns:
+            object(s) selected by `output` (cf. `simulate`), with the batch
+            axes after the ADC axis
+        """
+        options = self._options(kwargs)
+        variables, statics = self._bind(args, kwargs)
+        return self._signal(variables, statics, sample, scanner, options)
+
+    def build(self, *args, **kwargs):
+        """layer-2 sequence for given values (without batch dimensions)"""
+        variables, statics = self._bind(args, kwargs)
+        for name, value in variables.items():
+            array = value.values if is_axis(value) else value
+            if jnp.ndim(array) != len(self.variables[name]):
+                raise ValueError(f"build: variable {name!r} has batch dimensions")
+        values = {name: static.value for name, static in statics}
+        values.update({n: v.values if is_axis(v) else v for n, v in variables.items()})
+        return self._build(values)
+
+    def trace(self, *args, sample=None, scanner=None, select=None, **kwargs):
+        """states after every operator, for one element of the batch
+
+        The element of the variables' batch axes is selected, the sequence
+        built with core-shaped values, then traced (cf. `functions.trace`).
+        Usable with `plotting.plot_epg` (directly, or `plot_epg(seq, ...)`).
+
+        Args:
+            args, kwargs: arguments of the function and trace options (method,
+                init, dtype, dims, alignment, max_unroll, jit, disp)
+            sample, scanner: cf. `simulate` (optional)
+            select: element of the batch axes, as a dict {axis: index}: axes
+                of the variables (names, or batch_i for unnamed ones) and of
+                the sample/scanner; other axes: index 0
+
+        Returns:
+            Trace
+        """
+        given = {k: kwargs.pop(k) for k in list(kwargs) if k in OPTIONS}
+        invalid = set(given) - set(TRACE_OPTIONS)
+        if invalid:
+            raise ValueError(f"Option(s) not applicable to trace: {sorted(invalid)}")
+        options = {k: v for k, v in self.options.items() if k in TRACE_OPTIONS}
+        options.update(given)
+        variables, statics = self._bind(args, kwargs)
+
+        # selected element of the variables' batch axes
+        order, sizes, labels, arrays = _batch(variables, self.variables)
+        select = dict(select or {})
+        index = {}
+        for label in order:
+            i = select.pop(label, 0)
+            if not -sizes[label] <= i < sizes[label]:
+                raise IndexError(
+                    f"Index {i} out of range for axis {label!r} of size {sizes[label]}"
+                )
+            index[label] = i
+        values = {name: static.value for name, static in statics}
+        for name, array in arrays.items():
+            values[name] = array[tuple(index[label] for label in labels[name])]
+        return functions.trace(
+            self._build(values), sample, scanner, select=select or None, **options
+        )
+
+    def derivatives(
+        self,
+        *args,
+        jacobian=(),
+        hessian=None,
+        coupled=(),
+        sample=None,
+        scanner=None,
+        **kwargs,
+    ):
+        """signal, Jacobian and cross Hessian by name (cf. `diff.derivatives`)
+
+        Parameters are declared variables (batch dimensions: one column, core
+        dimensions: one column per value) or fields of the sample and scanner
+        (PD, T1, T2, df, B1, B0).
+
+        Args:
+            args, kwargs: arguments and options, as for `signal`
+            jacobian: names of the Jacobian's parameters
+            hessian: (rows, cols): names of the Hessian's rows and columns
+            coupled: names of sample/scanner axes along which elements
+                interact (e.g. compartments summed by the ADC)
+            sample, scanner: cf. `simulate`
+
+        Returns:
+            diff.Derivatives (value, jacobian (*shape, ncols), hessian
+            (*shape, nrows, ncols), labels)
+        """
+        options = self._options(kwargs)
+        options["output"] = "jax"
+        variables, statics = self._bind(args, kwargs)
+        values, cores, rebuild = self._parameters(variables, sample, scanner)
+
+        def signal(params):
+            variables_, sample_, scanner_ = rebuild(params)
+            return self._signal(variables_, statics, sample_, scanner_, dict(options))
+
+        return diff.derivatives(
+            signal,
+            values,
+            jacobian=jacobian,
+            hessian=hessian,
+            coupled=tuple(coupled) + cores,
+        )
+
+    def jacobian(self, names, *args, **kwargs):
+        """signal and Jacobian (*shape, ncols) by name (cf. `derivatives`)"""
+        d = self.derivatives(*args, jacobian=names, **kwargs)
+        return d.value, d.jacobian
+
+    def hessian(self, rows, cols=None, *args, **kwargs):
+        """signal, Jacobian (rows) and Hessian (*shape, nrows, ncols) by name"""
+        cols = rows if cols is None else cols
+        d = self.derivatives(*args, jacobian=rows, hessian=(rows, cols), **kwargs)
+        return d.value, d.jacobian, d.hessian
+
+    def crlb(
+        self,
+        names,
+        *args,
+        W=None,
+        sigma2=1,
+        log=False,
+        split=False,
+        axis="adc",
+        **kwargs,
+    ):
+        """Cramer-Rao lower bound of parameters (cf. `stats.crlb`)
+
+        Args:
+            names: parameters (declared variables, sample/scanner fields)
+            args, kwargs: arguments and options, as for `derivatives`
+            W, sigma2, log: cf. `stats.crlb`
+            split: return the bound of each parameter (`stats.crlb_split`)
+            axis: name(s) of the observation axes (default: the ADC axis),
+                e.g. ("adc", "point") when points are observations too
+
+        Returns:
+            bound for the remaining axes (*batch) (split: (*batch, nparam))
+        """
+        dims = self._dims(args, kwargs)
+        J = self.derivatives(*args, jacobian=names, **kwargs).jacobian
+        J = _fold(J, dims, axis, ntrail=1)
+        W = None if W is None else jnp.asarray(W)
+        bound = stats.crlb_split if split else stats.crlb
+        return bound(J, W=W, sigma2=sigma2, log=log)
+
+    def confint(
+        self, obs, names, *args, conflevel=0.95, hessian=False, axis="adc", **kwargs
+    ):
+        """confidence intervals of parameters (cf. `stats.confint`)
+
+        Args:
+            obs: observed signal (same shape as the simulated signal)
+            names: parameters (declared variables, sample/scanner fields)
+            args, kwargs: arguments (estimated values) and options
+            conflevel: confidence level
+            hessian: include the Hessian of the prediction
+            axis: name(s) of the observation axes (default: the ADC axis)
+
+        Returns:
+            intervals (*batch, nparam): half-widths, band (shape of `obs`):
+            half-widths of the confidence band of the prediction
+        """
+        dims = self._dims(args, kwargs)
+        H = (names, names) if hessian else None
+        d = self.derivatives(*args, jacobian=names, hessian=H, **kwargs)
+        obs = jnp.asarray(obs)
+        if obs.shape != d.value.shape:
+            raise ValueError(
+                f"Observation shape {obs.shape} differs from the prediction "
+                f"shape {d.value.shape}"
+            )
+        intervals, band = stats.confint(
+            _fold(obs, dims, axis),
+            _fold(d.value, dims, axis),
+            _fold(d.jacobian, dims, axis, ntrail=1),
+            None if d.hessian is None else _fold(d.hessian, dims, axis, ntrail=2),
+            conflevel=conflevel,
+        )
+        return intervals, _unfold(band, dims, axis, d.value.shape)
+
+    #
+    # private
+
+    def _options(self, kwargs):
+        """simulation options: decorator defaults, then call (popped from kwargs)"""
+        options = {**self.options}
+        options.update({k: kwargs.pop(k) for k in list(kwargs) if k in OPTIONS})
+        return options
+
+    def _parameters(self, variables, sample, scanner):
+        """differentiable values by name, coupled core axes, rebuilding function"""
+        values, cores, forms = {}, (), {}
+        for name, value in variables.items():
+            names = value.names if is_axis(value) else ()
+            array = value.values if is_axis(value) else value
+            nbatch = jnp.ndim(array) - len(self.variables[name])
+            batch = names or tuple(f"_{name}_batch{i}" for i in range(nbatch))
+            core = tuple(f"_{name}_core{i}" for i in range(len(self.variables[name])))
+            cores += core
+            forms[name] = names
+            values[name] = Axis(batch + core, array) if jnp.ndim(array) else array
+        objects = {"sample": sample, "scanner": scanner}
+        fields = {"sample": Sample.FIELDS, "scanner": ("B1", "B0")}
+        for key, obj in objects.items():
+            for field in fields[key] if obj is not None else ():
+                if getattr(obj, field) is not None:
+                    values[field] = getattr(obj, field)
+
+        def rebuild(params):
+            variables_ = {}
+            for name in variables:
+                array = params[name].values if is_axis(params[name]) else params[name]
+                variables_[name] = Axis(forms[name], array) if forms[name] else array
+            rebuilt = []
+            for key, obj in objects.items():
+                changes = {f: params[f] for f in fields[key] if f in params}
+                rebuilt.append(obj.replace(**changes) if obj is not None else None)
+            return (variables_, *rebuilt)
+
+        return values, cores, rebuild
+
+    def _dims(self, args, kwargs):
+        """axis names of the signal (single stacked output)"""
+        kwargs = {k: v for k, v in kwargs.items() if k != "coupled"}
+        layout = self.signal(*args, **{**kwargs, "output": "layout"})
+        if not isinstance(layout, OutputLayout):
+            raise ValueError(
+                "Statistics need a single stacked signal (no group=, ADCs of the "
+                "same shape)"
+            )
+        return layout.dims
+
+    def _bind(self, args, kwargs):
+        """variables (checked) and static arguments (hashable)"""
+        try:
+            bound = self.signature.bind(*args, **kwargs)
+        except TypeError as exc:
+            raise TypeError(f"{self.__name__}: {exc}") from None
+        bound.apply_defaults()
+        variables, statics = {}, []
+        for name, value in bound.arguments.items():
+            if name in self.variables:
+                variables[name] = self._check_variable(name, value)
+            else:
+                statics.append((name, Static(name, value)))
+        self._count_recompilations(statics)
+        return variables, tuple(statics)
+
+    def _check_variable(self, name, value):
+        """value as an array (or Axis) with its declared core dimensions"""
+        dims = self.variables[name]
+        names = value.names if is_axis(value) else ()
+        array = jnp.asarray(value.values if is_axis(value) else value)
+        if not jnp.issubdtype(array.dtype, jnp.inexact):
+            array = array.astype(jnp.result_type(float))
+        nbatch = array.ndim - len(dims)
+        declaration = _declaration(name, dims)
+        if nbatch < 0:
+            raise ValueError(
+                f"Variable {name!r} is declared {declaration} ({len(dims)} core "
+                f"dimension(s)), got shape {array.shape}"
+            )
+        elif names and nbatch != len(names):
+            raise ValueError(
+                f"Axis {value.name!r} of variable {name!r} names {len(names)} "
+                f"dimension(s), but the value {array.shape} has {nbatch} batch "
+                f"dimension(s) beyond the declared core ({declaration})"
+            )
+        return Axis(names, array) if names else array
+
+    def _count_recompilations(self, statics):
+        for name, static in statics:
+            if not static.array:
+                continue
+            seen = self._seen.setdefault(name, set())
+            if static in seen:
+                continue
+            seen.add(static)
+            if len(seen) == RECOMPILATIONS:
+                warnings.warn(
+                    f"Argument {name!r} of {self.__name__} changed {RECOMPILATIONS} "
+                    "times, each causing a recompilation: declare it as a variable "
+                    f"({name + '[]'!r}) if its values change between calls",
+                    stacklevel=5,
+                )
+
+    def _build(self, values):
+        sequence = self.function(**values)
+        if not isinstance(sequence, list):
+            raise TypeError(
+                f"{self.__name__} must return a list of operators, "
+                f"got {type(sequence).__name__}"
+            )
+        return sequence
+
+    def _signal(self, variables, statics, sample, scanner, options):
+        output = options.pop("output", "jax")
+        names = (output,) if isinstance(output, str) else tuple(output)
+        if not set(names) <= set(OUTPUTS):
+            raise ValueError(f"Unknown output(s): {set(names) - set(OUTPUTS)}")
+        jit = options.pop("jit", True)
+        dynamic = {"init": options.pop("init", None), "probe": None}
+        if not _is_static_probe(options.get("probe")):
+            dynamic["probe"] = options.pop("probe")
+        options = tuple(sorted((k, _hashable(v)) for k, v in options.items()))
+        self._check_names(variables, sample, scanner)
+
+        # values and times: compiled program
+        kinds = tuple(k for k in ("jax", "times") if _needs(k, names))
+        results = {}
+        if kinds:
+            run = self._compiled if jit else self._execute
+            out = run(
+                variables,
+                sample,
+                scanner,
+                dynamic,
+                statics=statics,
+                options=options,
+                kinds=kinds,
+            )
+            results.update(zip(kinds, _plain(out)))
+        if "layout" in names or "xarray" in names:
+            results["layout"] = self._layout(
+                variables, statics, sample, scanner, dynamic, options
+            )
+
+        group = dict(options).get("group")
+        outputs = []
+        for name in names:
+            if name == "numpy":
+                outputs.append(_to_numpy(results["jax"]))
+            elif name == "xarray":
+                outputs.append(_xarray(results, group))
+            else:
+                outputs.append(results[name])
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+    def _execute(self, variables, sample, scanner, dynamic, *, statics, options, kinds):
+        """values (and times) of the batch, in the compiled program"""
+        options = dict(options)
+        if dynamic["probe"] is not None:
+            options["probe"] = dynamic["probe"]
+        values = {name: static.value for name, static in statics}
+        order, _, labels, arrays = _batch(variables, self.variables)
+        masks = []
+
+        def core(vals):
+            sequence = self._build({**values, **vals})
+            out = simulate(
+                sequence,
+                sample,
+                scanner,
+                init=dynamic["init"],
+                output=kinds,
+                jit=False,
+                **options,
+            )
+            out = _ordered((out,) if len(kinds) == 1 else out)
+            masks.append(_adc_mask(out))  # structure, at trace time
+            return out
+
+        # map over the batch dimensions: outer axes first
+        mapped = core
+        for k in reversed(range(len(order))):
+            in_axes = {}
+            for name in arrays:
+                remaining = [lab for lab in labels[name] if lab not in order[:k]]
+                in_axes[name] = (
+                    remaining.index(order[k]) if order[k] in remaining else None
+                )
+            mapped = jax.vmap(mapped, in_axes=(in_axes,))
+        out = mapped(arrays)
+
+        # batch axes after the ADC axis (values acquired by stacked ADCs)
+        nbatch = len(order)
+        if nbatch:
+            out = jax.tree.map(
+                lambda leaf, adc: _after_adc(leaf, nbatch) if adc else leaf,
+                out,
+                masks[0],
+            )
+        return out
+
+    def _layout(self, variables, statics, sample, scanner, dynamic, options):
+        """layout(s) of the values: one simulation's layout, with the batch axes"""
+        order, sizes, labels, arrays = _batch(variables, self.variables)
+        values = {name: static.value for name, static in statics}
+        values.update({name: arrays[name][(0,) * len(labels[name])] for name in arrays})
+        options = dict(options)
+        if dynamic["probe"] is not None:
+            options["probe"] = dynamic["probe"]
+        options.pop("disp", None)
+        layout = simulate(
+            self._build(values),
+            sample,
+            scanner,
+            init=dynamic["init"],
+            output="layout",
+            **options,
+        )
+        return _insert_batch(layout, tuple(order), tuple(sizes[o] for o in order))
+
+    def _check_names(self, variables, sample, scanner):
+        """batch axes of the variables are not axes of the sample/scanner"""
+        named = {n for v in variables.values() if is_axis(v) for n in v.names}
+        layout = axes.get_layout(sample, scanner)
+        shared = named & set(layout.dims)
+        if shared:
+            raise ValueError(
+                f"Axis name(s) {sorted(shared)} used both by a variable and by the "
+                "sample/scanner (pairing across them is not supported)"
+            )
+
+
+class Static:
+    """Hashable key of a static argument (the function receives the value)
+
+    Attributes:
+        value: the argument
+        key: hashable key (arrays: shape, dtype and bytes)
+        array: True for array-like values (recompilation warning)
+    """
+
+    def __init__(self, name, value):
+        self.value = value
+        self.key = _static_key(name, value)
+        self.array = _is_array_like(value)
+
+    def __hash__(self):
+        return hash(self.key)
+
+    def __eq__(self, other):
+        return isinstance(other, Static) and self.key == other.key
+
+    def __repr__(self):
+        return f"Static({self.value!r})"
+
+
+#
+# functions
+
+
+def _declaration(name, dims):
+    if not dims:
+        return name
+    return f"{name}[{', '.join(d for d in dims if d != f'{name}[]')}]"
+
+
+def _static_key(name, value):
+    if isinstance(value, jax.core.Tracer):
+        raise TypeError(
+            f"Argument {name!r} is static but traced (e.g. under jax.grad or "
+            "jax.jit): declare it as a variable"
+        )
+    elif _is_array_like(value):
+        # numeric arrays, lists and tuples: same key for the same contents
+        array = np.asarray(value)
+        return (np.ndarray, array.shape, array.dtype.str, array.tobytes())
+    elif isinstance(value, (list, tuple)):
+        return (type(value), tuple(_static_key(name, v) for v in value))
+    elif isinstance(value, dict):
+        return (
+            dict,
+            tuple(sorted((k, _static_key(name, v)) for k, v in value.items())),
+        )
+    try:
+        hash(value)
+    except TypeError:
+        raise TypeError(
+            f"Static argument {name!r} is not hashable ({type(value).__name__}): "
+            "pass a tuple or an array, or declare it as a variable"
+        ) from None
+    return (type(value), value)
+
+
+def _is_array_like(value):
+    if isinstance(value, (np.ndarray, jax.Array)):
+        return True
+    elif isinstance(value, (list, tuple)) and value:
+        try:
+            return np.asarray(value).dtype.kind in "iufc"
+        except ValueError:
+            return False
+    return False
+
+
+def _hashable(value):
+    """options as static arguments of jit (lists to tuples)"""
+    if isinstance(value, list):
+        return tuple(_hashable(v) for v in value)
+    return value
+
+
+def _is_static_probe(probe):
+    """attribute names (static) vs Probe operators (pytrees)"""
+    if isinstance(probe, tuple):
+        return all(isinstance(p, str) for p in probe)
+    return probe is None or isinstance(probe, str)
+
+
+def _needs(kind, names):
+    if kind == "jax":
+        return bool({"jax", "numpy", "xarray"} & set(names))
+    return bool({"times", "xarray"} & set(names))
+
+
+def _batch(variables, cores):
+    """batch structure of the variables
+
+    Returns:
+        order: batch axis labels, in order of first appearance
+        sizes: dict label -> size
+        labels: dict variable -> labels of its leading batch dimensions
+        arrays: dict variable -> array (unnamed batch dimensions broadcast)
+    """
+    labels, arrays, unnamed = {}, {}, {}
+    for name, value in variables.items():
+        array = value.values if is_axis(value) else value
+        nbatch = jnp.ndim(array) - len(cores[name])
+        arrays[name] = array
+        if is_axis(value):
+            labels[name] = value.names
+        elif nbatch:
+            unnamed[name] = jnp.shape(array)[:nbatch]
+        else:
+            labels[name] = ()
+    if unnamed:
+        try:
+            common = tuple(np.broadcast_shapes(*unnamed.values()))
+        except ValueError:
+            raise ValueError(f"Incompatible unnamed batch shapes: {unnamed}") from None
+        for name, shape in unnamed.items():
+            core = jnp.shape(arrays[name])[len(shape) :]
+            arrays[name] = jnp.broadcast_to(arrays[name], common + core)
+            labels[name] = tuple(f"batch_{i}" for i in range(len(common)))
+
+    order, sizes, core_sizes = [], {}, {}
+    for name in variables:
+        shape = jnp.shape(arrays[name])
+        for label, size in zip(labels[name], shape):
+            if sizes.setdefault(label, size) != size:
+                raise ValueError(
+                    f"Inconsistent sizes for batch axis {label!r}: "
+                    f"{sizes[label]}, {size} ({name!r})"
+                )
+            if label not in order:
+                order.append(label)
+        for dim, size in zip(cores[name], shape[len(labels[name]) :]):
+            if core_sizes.setdefault(dim, size) != size:
+                raise ValueError(
+                    f"Inconsistent sizes for core dimension {dim!r}: "
+                    f"{core_sizes[dim]}, {size} ({name!r})"
+                )
+    return order, sizes, labels, arrays
+
+
+def _fold(array, dims, axis, ntrail=0):
+    """observation axes (names) first, folded into one (statistics layout)
+
+    Args:
+        array: (*shape, *trail) with axis names `dims` for `shape`
+        axis: name or tuple of names of the observation axes
+        ntrail: number of trailing axes (parameters)
+    """
+    index = _axis_index(dims, axis)
+    moved = jnp.moveaxis(array, index, tuple(range(len(index))))
+    nobs = int(np.prod(moved.shape[: len(index)]))
+    return moved.reshape((nobs,) + moved.shape[len(index) :])
+
+
+def _unfold(array, dims, axis, shape):
+    """inverse of `_fold` (no trailing axes), back to `shape`"""
+    index = _axis_index(dims, axis)
+    array = array.reshape(tuple(shape[i] for i in index) + array.shape[1:])
+    return jnp.moveaxis(array, tuple(range(len(index))), index)
+
+
+def _axis_index(dims, axis):
+    axis = (axis,) if isinstance(axis, str) else tuple(axis)
+    for name in axis:
+        if name not in dims:
+            raise ValueError(f"Unknown observation axis {name!r}, available: {dims}")
+    return tuple(dims.index(name) for name in axis)
+
+
+def _adc_mask(tree):
+    """True for leaves with a leading ADC axis (stacked values)"""
+    if isinstance(tree, list):
+        return jax.tree.map(lambda _: False, tree)
+    elif isinstance(tree, tuple):
+        return tuple(_adc_mask(item) for item in tree)
+    elif isinstance(tree, dict):
+        return type(tree)((key, _adc_mask(item)) for key, item in tree.items())
+    return True
+
+
+def _ordered(tree):
+    """dicts as OrderedDict (JAX sorts the keys of plain dicts, e.g. groups)"""
+    if isinstance(tree, dict):
+        return collections.OrderedDict((k, _ordered(v)) for k, v in tree.items())
+    elif isinstance(tree, (list, tuple)):
+        return type(tree)(_ordered(item) for item in tree)
+    return tree
+
+
+def _plain(tree):
+    """OrderedDict back to plain dicts (order kept)"""
+    if isinstance(tree, dict):
+        return {k: _plain(v) for k, v in tree.items()}
+    elif isinstance(tree, (list, tuple)):
+        return type(tree)(_plain(item) for item in tree)
+    return tree
+
+
+def _after_adc(leaf, nbatch):
+    """move the leading batch axes after the ADC axis"""
+    return jnp.moveaxis(leaf, tuple(range(nbatch)), tuple(range(1, nbatch + 1)))
+
+
+def _insert_batch(tree, dims, shape):
+    """layout(s) with the batch axes (after the ADC axis, if any)"""
+    if isinstance(tree, OutputLayout):
+        n = 1 if tree.dims[:1] == ("adc",) else 0
+        return OutputLayout(
+            tree.dims[:n] + dims + tree.dims[n:],
+            tree.shape[:n] + shape + tree.shape[n:],
+            tree.coords,
+            tree.names,
+        )
+    elif isinstance(tree, dict):
+        return {key: _insert_batch(item, dims, shape) for key, item in tree.items()}
+    elif isinstance(tree, (list, tuple)):
+        return type(tree)(_insert_batch(item, dims, shape) for item in tree)
+    return tree
+
+
+def _xarray(results, group):
+    """xarray output (ADC times as coordinates if the same for the whole batch)"""
+    values, layout, times = results["jax"], results["layout"], results["times"]
+    if group is None:
+        return _to_xarray(values, layout, _common_times(times))
+    return {
+        key: _to_xarray(values[key], layout[key], _common_times(times[key]))
+        for key in values
+    }
+
+
+def _common_times(times):
+    """ADC times (nadc,) if the same for every batch element, else None"""
+    times = np.asarray(_to_numpy(times))
+    flat = times.reshape(len(times), -1)
+    return flat[:, 0] if np.allclose(flat, flat[:, :1]) else None