Program combinators

Program combinators.

The implement is pretty much backend-agnostic. We just assume that the core backend supports the following functionality:

  • suffix(p): rename latent variables of the program p,

  • traced_evaluate(p, latents=None): execute p and collect trace, metrics, optionally we can substitute values in latents to p,

  • empirical(out, trace, metrics): create a delta program given output, trace, and metrics. Inputs of empirical are outputs of traced_evaluate.

compose(q2, q1, suffix=True)[source]

Executes q2(*q1(…)).

Note: We only allow at most one of q1 or q2 is weighted.

Parameters:

  • q2 – a program

  • q1 – a program

  • suffix – whether to add suffix _PREV_ to variables in q1

Returns:

the composed program

Return type:

q

extend(p, f)[source]

Executes f(*p(…)) with random variables in f marked as auxiliary.

Note: We don’t allow recursively marginalize out p yet.

Parameters:

  • p – a target program

  • f – an auxiliary program

Returns:

the extended program

Return type:

p_new

fori_loop(lower, upper, body_fun, init_program)[source]

Returns a program which loops over programs created by body_fun.

Parameters:

  • lower – loop index lower bound

  • upper – loop index upper bound (exclusive)

  • body_fun – a function that takes a pair of inputs (index, program) and return a new program

  • init_program – initial program for body_fun

Returns:

the final program

Return type:

q

propose(p, q, *, loss_fn=None, detach=False, chain=False)[source]

Returns a new program with important weight.

We assume the leftmost batch dimension is the particle dimension. You can add additional batch dimensions to the whole program by using vmap, e.g. vmap(propose(p, q)).

Note: We assume superfluous variables, which appear in q but not in p, implicitly follow Delta distribution in p.

Parameters:

  • p – a target program

  • q – a proposal program

  • loss_fn – a function that computes loss of this propose combinator

  • detach – whether to detach value of the returned program

  • chain – if True, we will use output of q as input of p

Returns:

the proposed program

Return type:

q_new

resample(q, num_samples=None)[source]

Returns a new program with equally-weighted particles.

Parameters:

  • q – a program

  • num_samples – the number of samples after resampling. Set this to an empty tuple to draw 1 sample without the leftmost singleton dimension. Defaults to the number of particles in q.

Returns:

the resampled program

Return type:

q_new