Core / factor graphs

The directed factor graph beneath the circuits: a graph of factors, each a sampler over its sites.

Abstract base classes

AbstractTiledFactorclass

AbstractTiledFactor(
    base: AbstractFactor,
    n_tiles: int,
    weight_tied: bool,
    *,
    batch_size: int | None = None,
    slice_info: bool = False,
)

Base implementation for factors tiled across independent replicas.

batch_size chooses the mapping mode: None, 0, or any value >= n_tiles uses jax.vmap (all tiles at once); a smaller value uses jax.lax.map with chunks of that size.

The info argument is the runtime info forwarded to each tile's base.sample. By default the same info is broadcast to every tile; set slice_info=True to instead pass per-tile info.

Arguments:

base: The factor to replicate.
n_tiles: Number of tiles.
weight_tied: Whether all tiles share one parameter set (see above).
batch_size: vmap vs lax.map execution strategy (see above).
slice_info: Whether info is sliced per tile (see above).

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Sample every tile and stack the results along a leading n_tiles axis. With return_aux=True each tile's aux is stacked likewise.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Per-tile AbstractFactor.sample_with_references.

Each tile calls base.sample_with_references(n_references) and the results are stacked across tiles. The main output leaves have shape (n_tiles, *base_output_shape) and aux leaves have shape (n_references + 1, n_tiles, *base_aux_leaf_shape).

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

base.init_params when weight_tied, else vmap-ed over n_tiles.

AbstractChainFactorclass

AbstractChainFactor(
    base: AbstractFactor,
    n_steps: int,
    feedback_porting_fn: Union[Callable[[PyTree[Array]], Mapping[str, PyTree[Array]]], str],
    weight_tied: bool,
    *,
    slice_info: bool = False,
)

Base implementation for factors chained via jax.lax.scan.

At each step the base's input ports split into two disjoint sets:

Feedback ports: receive feedback_porting_fn(previous step's main output); at step 0 they take the chain's initial state.
Broadcast ports: get the caller's value unchanged at every step.

feedback_porting_fn is either a port-name str (for lambda main: {str: main}) or a Callable mapping the previous main output to a dict of feedback-port values.

Arguments:

base: The factor applied at every step.
n_steps: Number of steps.
feedback_porting_fn: Port-name str or Callable routing each step's main output into the next step.
weight_tied: Whether all steps share one parameter set.
slice_info: Whether info is sliced per step.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Run the scan and return the final step's output. With return_aux=True additionally returns the per-step aux stacked along a leading axis of size n_steps.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Per-step AbstractFactor.sample_with_references along the scan.

Each step calls base.sample_with_references(n_references); the main output of the step is threaded as the scan carry.The returned aux_trace leaves have shape (n_references + 1, n_steps, *base_aux_leaf_shape).

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

vmap(base.init_params) per step, or one shared init when weight_tied.

AbstractDFGclass

AbstractDFG()

A Factor built as a DAG of placed factors.

Holds the structure and concretises sample / sample_with_references as an eager topological walk.

The DFG owns a flat namespace of addresses, one per input port and one per Site, and the two sets must be disjoint. Every parents entry and output_name is such an address. A DFG is itself a Factor, so it can nest as the factor of a Site.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: DFGInfo | None = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], tuple[PyTree[Array], ...]]

Run the DAG once and return the output_name value.

Samples each site in topological order, routing parent outputs through its porting_fn.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: DFGInfo | None = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], tuple[PyTree[Array], ...]]

Like sample, but draws each site via its factor's sample_with_references, so every site yields n_references + 1 samples. Returns (output, per_site_auxes).

distribute_paramsmethod

distribute_params(params: PyTree[Array]) -> tuple[PyTree[Array], ...]

Scatter the shared params mapping into a per-site tuple.

gather_param_gradsmethod

gather_param_grads(
    params: PyTree[Array],
    site_grads: Mapping[int, PyTree[Array]],
) -> PyTree[Array]

Gather per-site parameter gradients into the shared params dict.

distribute_infomethod

distribute_info(info: DFGInfo | None) -> tuple[PyTree, ...]

Scatter info.entries into a per-site tuple.

AbstractFactorclass

AbstractFactor()

Base class for probabilistic factors.

A directed factor is a conditional distribution $P(\text{output} \mid \text{inputs})$.

Factor methods operate with the following:

params: the factor's parameter pytree.
info: optional runtime auxiliary info.
site_info: optional per-site metadata supplied by the surrounding Site.
return_aux: when True, methods additionally return a factor-defined aux pytree.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Draw a "main" sample plus n_references additional samples.

Always returns (main_output, aux). main_output is a single sample. Every aux leaf gains a leading axis of size n_references + 1: position 0 is the main sample's aux, positions 1 .. n_references are the references' auxes.

For a sample with no references and no aux, use sample.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_references: Number of reference samples to draw alongside the main sample.

Returns:

(main_output, aux).

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

AbstractReferenceFactorclass

AbstractReferenceFactor()

Abstract factor + a concrete sample_with_references.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Generic reference-sampling default.

Draws n_references + 1 samples via sample_multiple and returns the first as the main sample, with all of their auxes stacked.

AbstractHasLogProbabilityclass

AbstractHasLogProbability()

Capability mixin for factors with a tractable, analytic log_probability(outputs | inputs).

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

log_probabilitymethod

log_probability(
    inputs: Mapping[str, PyTree[Array]],
    outputs: PyTree[Array],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> Float[Array, ''] | tuple[Float[Array, ''], PyTree[Array]]

Scalar log-probability of outputs given inputs.

Arguments:

inputs: Per-port pytree inputs, as in sample.
outputs: An output pytree matching output_spec.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
return_aux: Whether to additionally return a factor-defined aux pytree.

Returns:

The scalar log P(outputs | inputs), or (log_p, aux) when return_aux=True.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Draw a "main" sample plus n_references additional samples.

For a sample with no references and no aux, use sample.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_references: Number of reference samples to draw alongside the main sample.

Returns:

(main_output, aux).

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

AbstractEnumerableOutputFactorclass

AbstractEnumerableOutputFactor()

Factor whose output state space is finite, with a canonical ordering.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Draw a "main" sample plus n_references additional samples.

For a sample with no references and no aux, use sample.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_references: Number of reference samples to draw alongside the main sample.

Returns:

(main_output, aux).

n_output_statesproperty

n_output_states

Number of output states in the canonical ordering.

get_nth_output_statemethod

get_nth_output_state(n: int | Int[Array, '']) -> PyTree[Array]

Return the n-th output state in the canonical ordering.

The returned pytree has the same structure as the output of AbstractFactor.sample.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

output_state_to_indexmethod

output_state_to_index(outputs: PyTree[Array]) -> Float[Array, '']

Index of outputs in the canonical ordering, as a scalar float.

AbstractFiniteStateSpaceFactorclass

AbstractFiniteStateSpaceFactor()

Factor whose input and output state spaces are both finite.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Draw a "main" sample plus n_references additional samples.

For a sample with no references and no aux, use sample.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_references: Number of reference samples to draw alongside the main sample.

Returns:

(main_output, aux).

n_output_statesproperty

n_output_states

Number of output states in the canonical ordering.

get_nth_output_statemethod

get_nth_output_state(n: int | Int[Array, '']) -> PyTree[Array]

Return the n-th output state in the canonical ordering.

The returned pytree has the same structure as the output of AbstractFactor.sample.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

output_state_to_indexmethod

output_state_to_index(outputs: PyTree[Array]) -> Float[Array, '']

Index of outputs in the canonical ordering, as a scalar float.

n_input_statesproperty

n_input_states

Number of input states in the canonical ordering.

get_nth_input_statemethod

get_nth_input_state(n: int | Int[Array, '']) -> Mapping[str, PyTree[Array]]

Return the n-th input state in the canonical ordering.

The returned pytree has the same structure as the inputs argument to AbstractFactor.sample.

input_state_to_indexmethod

input_state_to_index(inputs: Mapping[str, PyTree[Array]]) -> Float[Array, '']

Index of inputs in the canonical ordering, as a scalar float.

AbstractHasExplicitOutputDistributionclass

AbstractHasExplicitOutputDistribution()

Factor whose conditional is available in closed form as an explicit log-prob vector over the enumerable output states.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Draw a "main" sample plus n_references additional samples.

For a sample with no references and no aux, use sample.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_references: Number of reference samples to draw alongside the main sample.

Returns:

(main_output, aux).

n_output_statesproperty

n_output_states

Number of output states in the canonical ordering.

get_nth_output_statemethod

get_nth_output_state(n: int | Int[Array, '']) -> PyTree[Array]

Return the n-th output state in the canonical ordering.

The returned pytree has the same structure as the output of AbstractFactor.sample.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

output_state_to_indexmethod

output_state_to_index(outputs: PyTree[Array]) -> Float[Array, '']

Index of outputs in the canonical ordering, as a scalar float.

get_log_output_distributionmethod

get_log_output_distribution(
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> Float[Array, 'n_output_states'] | tuple[Float[Array, 'n_output_states'], PyTree[Array]]

Log-probabilities of all n_output_states outputs given inputs.

Returns a vector v of shape (n_output_states,) with v[j] = log P(get_nth_output_state(j) | inputs).

Arguments:

inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
return_aux: Whether to additionally return a factor-defined aux.

Returns:

The length-n_output_states log-probability vector, or (v, aux) when return_aux=True.

log_probabilitymethod

log_probability(
    inputs: Mapping[str, PyTree[Array]],
    outputs: PyTree[Array],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> Float[Array, ''] | tuple[Float[Array, ''], PyTree[Array]]

Index get_log_output_distribution at the queried outputs.

AbstractMatrixFactorclass

AbstractMatrixFactor()

A finite-state factor whose conditional is given explicitly as a matrix.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw a single sample from this factor.

Arguments:

key: PRNG key.
inputs: Per-port pytree inputs. Keys must equal the self.input_ports keys.
params: Parameter pytree for this factor.
info: Runtime auxiliary info.
site_info: Static per-site metadata supplied by the surrounding Site.
return_aux: If True, return (output, aux).

Returns:

The sampled output (a pytree matching output_spec), or (output, aux) when return_aux=True.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Generic reference-sampling default.

Draws n_references + 1 samples via sample_multiple and returns the first as the main sample, with all of their auxes stacked.

log_probabilitymethod

log_probability(
    inputs: Mapping[str, PyTree[Array]],
    outputs: PyTree[Array],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> Float[Array, ''] | tuple[Float[Array, ''], PyTree[Array]]

Index get_log_output_distribution at the queried outputs.

input_portsproperty

input_ports

output_specproperty

output_spec

n_input_statesproperty

n_input_states

n_output_statesproperty

n_output_states

get_nth_input_statemethod

get_nth_input_state(n: int | Int[Array, '']) -> Mapping[str, PyTree[Array]]

See AbstractFiniteStateSpaceFactor.get_nth_input_state.

get_nth_output_statemethod

get_nth_output_state(n: int | Int[Array, '']) -> PyTree[Array]

See AbstractFiniteStateSpaceFactor.get_nth_output_state.

input_state_to_indexmethod

input_state_to_index(inputs: Mapping[str, PyTree[Array]]) -> Float[Array, '']

See AbstractFiniteStateSpaceFactor.input_state_to_index.

output_state_to_indexmethod

output_state_to_index(outputs: PyTree[Array]) -> Float[Array, '']

See AbstractEnumerableOutputFactor.output_state_to_index.

get_log_probability_matrixmethod

get_log_probability_matrix(
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
) -> Float[Array, 'n_input_states n_output_states']

Return the (n_input_states, n_output_states) log-probability matrix.

Entry [i, j] is log P(output_states[j] | input_states[i]).

Arguments:

params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.

Returns:

The matrix of log-probabilities.

get_log_output_distributionmethod

get_log_output_distribution(
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> Float[Array, 'n_output_states'] | tuple[Float[Array, 'n_output_states'], PyTree[Array]]

The matrix row selected by the input's canonical index.

Concrete classes

TiledFactorclass

TiledFactor(
    base: AbstractFactor,
    n_tiles: int,
    weight_tied: bool,
    *,
    batch_size: int | None = None,
    slice_info: bool = False,
)

Replication of a base factor across n_tiles independent tiles.

baseattribute

base: AbstractFactor

n_tilesattribute

n_tiles: int

weight_tiedattribute

weight_tied: bool

batch_sizeattribute

batch_size: int | None

slice_infoattribute

slice_info: bool

input_portsattribute

input_ports: Mapping[str, jaxPyTree[jax.ShapeDtypeStruct]]

output_specattribute

output_spec: jaxPyTree[jax.ShapeDtypeStruct]

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Sample every tile and stack the results along a leading n_tiles axis. With return_aux=True each tile's aux is stacked likewise.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Per-tile AbstractFactor.sample_with_references.

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

base.init_params when weight_tied, else vmap-ed over n_tiles.

ChainFactorclass

ChainFactor(
    base: AbstractFactor,
    n_steps: int,
    feedback_porting_fn: Union[Callable[[PyTree[Array]], Mapping[str, PyTree[Array]]], str],
    weight_tied: bool,
    *,
    slice_info: bool = False,
)

Sequential composition of a base factor via jax.lax.scan.

baseattribute

base: AbstractFactor

n_stepsattribute

n_steps: int

feedback_porting_fnattribute

feedback_porting_fn: Optional[Callable[[jaxPyTree[Array]], Mapping[str, jaxPyTree[Array]]]]

feedback_portsattribute

feedback_ports: tuple[str, ...]

all_step_input_portsattribute

all_step_input_ports: tuple[str, ...]

weight_tiedattribute

weight_tied: bool

slice_infoattribute

slice_info: bool

input_portsattribute

input_ports: Mapping[str, jaxPyTree[jax.ShapeDtypeStruct]]

output_specattribute

output_spec: jaxPyTree[jax.ShapeDtypeStruct]

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Run the scan and return the final step's output. With return_aux=True additionally returns the per-step aux stacked along a leading axis of size n_steps.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Per-step AbstractFactor.sample_with_references along the scan.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

vmap(base.init_params) per step, or one shared init when weight_tied.

Siteclass

Site(
    name: str,
    factor: AbstractFactor,
    parents: Any,
    porting_fn: Any,
    param_key: str | int | None,
    info_key: str | None,
    site_info: Any,
)

The placement of a Factor at one position in a DFG.

constructor

Construct a Site.

Arguments:

name: This site's address in the DFG namespace. Unique among sites and disjoint from the input-port names.
factor: The Factor placed here.
parents: Addresses (input ports or site names) feeding this factor; porting_fn maps them onto its named input ports.
porting_fn: How parents route into the factor's input dict: a tuple of port names (1:1 with parents), or a callable parents -> input dict for non-trivial routing.
param_key: Address of this site's parameters within the DFG's params.
info_key: Like param_key, for runtime info; None passes info=None.
site_info: per-site metadata passed to the factor.

nameattribute

name: str

factorattribute

factor: AbstractFactor

parentsattribute

parents: tuple[str, ...]

porting_fnattribute

porting_fn: Union[Callable[[Sequence[jaxPyTree[Array]]], Mapping[str, jaxPyTree[Array]]], tuple[str, ...]]

param_keyattribute

param_key: str | int | None

info_keyattribute

info_key: str | None

site_infoattribute

site_info: Any

DFGInfoclass

DFGInfo(
    expose_site_outputs: bool,
    entries: Mapping[str, PyTree[Array]] = <factory>,
)

DFG-level runtime info.

Separate from per-site info, which lives in entries under each site's info_key; this configures how the DAG itself runs.

Arguments:

expose_site_outputs: When True (and aux is requested), prepend a name-keyed dict of every site's main output to the aux return.
entries: Per-info_key mapping, scattered to sites by distribute_info. A child-DFG site's entry is itself a DFGInfo.

expose_site_outputsattribute

expose_site_outputs: bool

entriesattribute

entries: Mapping[str, jaxPyTree[Array]]

DFGclass

DFG(
    sites: tuple[Site, ...],
    input_ports: Mapping[str, PyTree[jax.ShapeDtypeStruct]],
    output_name: str,
)

A concrete, eagerly-walked DAG of factors.

sitesattribute

sites: tuple[Site, ...]

input_portsattribute

input_ports: Mapping[str, jaxPyTree[jax.ShapeDtypeStruct]]

output_specattribute

output_spec: jaxPyTree[jax.ShapeDtypeStruct]

output_nameattribute

output_name: str

topological_orderattribute

topological_order: tuple[int, ...]

sites_by_nameattribute

sites_by_name: Mapping[str, int]

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: DFGInfo | None = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], tuple[PyTree[Array], ...]]

Run the DAG once and return the output_name value.

Samples each site in topological order, routing parent outputs through its porting_fn.

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: DFGInfo | None = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], tuple[PyTree[Array], ...]]

Like sample, but draws each site via its factor's sample_with_references, so every site yields n_references + 1 samples. Returns (output, per_site_auxes).

distribute_paramsmethod

distribute_params(params: PyTree[Array]) -> tuple[PyTree[Array], ...]

Scatter the shared params mapping into a per-site tuple.

gather_param_gradsmethod

gather_param_grads(
    params: PyTree[Array],
    site_grads: Mapping[int, PyTree[Array]],
) -> PyTree[Array]

Gather per-site parameter gradients into the shared params dict.

distribute_infomethod

distribute_info(info: DFGInfo | None) -> tuple[PyTree, ...]

Scatter info.entries into a per-site tuple.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Initialise parameters once per distinct param_key.

DeterministicFactorclass

DeterministicFactor(
    fn: Callable[[Mapping[str, PyTree[Array]], Any], PyTree[Array]],
    input_ports: Any,
    output_spec: PyTree[jax.ShapeDtypeStruct],
)

A factor whose output is a deterministic function of its inputs.

Arguments:

fn: A pure function fn(inputs, site_info) -> output mapping the per-port inputs dict (and the surrounding Site's static site_info) to an output pytree matching output_spec.
input_ports: The factor's input-port specs.
output_spec: The spec of fn's output.

fnattribute

fn: Callable[[Mapping[str, jaxPyTree[Array]], Any], jaxPyTree[Array]]

input_portsattribute

input_ports: Mapping[str, jaxPyTree[jax.ShapeDtypeStruct]]

output_specattribute

output_spec: jaxPyTree[jax.ShapeDtypeStruct]

sample_multiplemethod

sample_multiple(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_samples: int = 1,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Draw n_samples samples by vmap-ing sample.

Arguments:

key: PRNG key, split n_samples ways.
inputs: Per-port pytree inputs, as in sample.
params: Parameter pytree for this factor.
info: Runtime auxiliary info, as in sample.
site_info: Static per-site metadata, as in sample.
n_samples: Number of independent samples to draw.
return_aux: Whether to return the aux pytree, as in sample.

Returns:

The stacked outputs, or (outputs, auxes) when return_aux=True, each with a leading axis of size n_samples.

samplemethod

sample(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> PyTree[Array] | tuple[PyTree[Array], PyTree[Array]]

Return fn(inputs, site_info).

sample_with_referencesmethod

sample_with_references(
    key: Key[Array, ''],
    inputs: Mapping[str, PyTree[Array]],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    n_references: int = 1,
) -> tuple[PyTree[Array], PyTree[Array]]

Evaluate fn once.

init_paramsmethod

init_params(key: Key[Array, '']) -> PyTree[Array]

Return a freshly-initialised parameter pytree.

Arguments:

key: PRNG key.

log_probabilitymethod

log_probability(
    inputs: Mapping[str, PyTree[Array]],
    outputs: PyTree[Array],
    params: PyTree[Array],
    info: PyTree = None,
    site_info: Any = None,
    return_aux: bool = False,
) -> Float[Array, ''] | tuple[Float[Array, ''], PyTree[Array]]

0 if outputs == fn(inputs, site_info), else -inf.