hyperiax.prebuilt.bffg.discrete_forward_sweep

hyperiax.prebuilt.bffg.discrete_forward_sweep(mean_fn, covar_fn)[source]

Build the unconditional forward-sampling down-sweep.

Draws each non-root state from the true 1-step Gaussian kernel \(X_t \mid X_{pa} = x \sim \mathcal N(\mu(x), Q(x))\) using the pre-stored noise \(z \sim \mathcal N(0, I)\) in node.z:

val = mean_fn(parent.val, params)
      + chol(covar_fn(parent.val, params)) @ node.z
Parameters:

  • mean_fn – Callable (x_parent, params) -> (d,) array returning the true conditional mean \(\mu(x)\).

  • covar_fn – Callable (x_parent, params) -> (d, d) SPD array returning the true conditional covariance \(Q(x)\).

Return type:

SweepFn

Returns:

A hyperiax.SweepFn that writes val at every non-root node. The root’s val must be set by the caller (typically via init_discrete_tree()’s root_val).