sweep (pmrf.sweep)

pmrf.sweep(eval_fn: Callable[[...], Any], *args: Any, grid: bool = False, template: Any = None) → Any

Vectorizes a function over the given input arguments.

Safely handles models as input by filtering out non-vmappable static fields (like strings, objects, and NumPy arrays).

Parameters:

• eval_fn (Callable) – The function to run the sweep over.

• *args (Any) – The arguments to evaluate. Can be raw arrays or full PyTree Models.

• grid (bool, default=False) – If False, performs a standard parallel sweep (zip-like) over the leading dimension of all dynamic arrays in args. If True, computes the cartesian product (meshgrid) of the arguments.

• template (Any, optional) – An unbatched PyTree (or tuple of PyTrees mirroring *args) used as a structural template to automatically extract batch axes. Required when sweeping over complex, pre-batched PyTrees (like Bayesian samples). Cannot be used concurrently with grid=True.

Returns:

The batched output of eval_fn.

Return type:

Any