vscan#

gpjax.scan.vscan(f, init, xs, length=None, reverse=False, unroll=1, log_rate=10, log_value=True)[source]#

Scan with verbose output.

This is based on code from this [excellent blog post](https://www.jeremiecoullon.com/2021/01/29/jax_progress_bar/).

Example

>>> import jax.numpy as jnp
...
>>> def f(carry, x):
...     return carry + x, carry + x
>>> init = 0
>>> xs = jnp.arange(10)
>>> vscan(f, init, xs)
(Array(45, dtype=int32), Array([ 0,  1,  3,  6, 10, 15, 21, 28, 36, 45], dtype=int32))
Parameters:

  • f (Callable[[Carry, X], Tuple[Carry, Y]]) – A function that takes in a carry and an input and returns a tuple of a new carry and an output.

  • init (Carry) – The initial carry.

  • xs (X) – The inputs.

  • length (Optional[int]) – The length of the inputs. If None, then the length of the inputs is inferred.

  • reverse (bool) – Whether to scan in reverse.

  • unroll (int) – The number of iterations to unroll.

  • log_rate (int) – The rate at which to log the progress bar.

  • log_value (bool) – Whether to log the value of the objective function.

Returns:

A tuple of the final carry and the outputs.

Return type:

Tuple[Carry, list[Y]]