sweep

Contents

sweep#

class braintools.fit.sweep(fn, grid, *, mode='product', batch_size=None)#

Evaluate a pure function over a grid of parameter values.

Parameters:

  • fn (Callable[[Dict[str, Any]], Any]) – fn(params) -> PyTree, a pure (vmap-able) function of a dict of scalar parameters.

  • grid (Dict[str, Any]) – {name: values} with 1-D arrays (or brainunit Quantities) of candidate values.

  • mode (str) – 'product' evaluates the cartesian product of the value arrays (grid shape (len(v1), len(v2), ...) in the order of grid); 'zip' pairs equal-length arrays element-wise (grid shape (n,)).

  • batch_size (int | None) – Number of grid points evaluated together by one vmap-ped call. The chunks run one after another through jax.lax.map(), so the peak memory is bounded by batch_size evaluations; a final partial chunk is handled when batch_size does not divide the grid size. None evaluates the whole grid in a single vmap.

Return type:

Tuple[Any, Dict[str, Any]]

Returns:

  • outputs (PyTree) – The outputs of fn stacked over the grid: every leaf has shape grid_shape + leaf_shape.

  • table (dict) – The flattened parameter table {name: values}; every entry has length prod(grid_shape) and row i is the input of output outputs.reshape(-1, ...)[i].

Raises:

ValueError – If grid is empty, a value array is not 1-D or is empty, mode is unknown, the arrays have different lengths in 'zip' mode, or batch_size is not a positive integer.

See also

Fitter

Optimise the parameters instead of scanning them.

Notes

The grid shape is static, so sweep can be called inside jax.jit; it is differentiable with respect to values fn closes over.

Examples

>>> import jax.numpy as jnp
>>> from braintools.fit import sweep
>>> out, table = sweep(lambda p: p['a'] * p['b'],
...                    {'a': jnp.array([1., 2.]), 'b': jnp.array([10., 20., 30.])})
>>> out
Array([[10., 20., 30.],
       [20., 40., 60.]], dtype=float32)
>>> table['b']
Array([10., 20., 30., 10., 20., 30.], dtype=float32)
>>>
>>> out, _ = sweep(lambda p: p['a'] + p['b'],
...                {'a': jnp.arange(5.), 'b': jnp.ones(5)}, mode='zip', batch_size=2)
>>> out
Array([1., 2., 3., 4., 5.], dtype=float32)
Contents