Circuits

A circuit is an ordered list of gates applied to an initial state.

Abstract base classes

AbstractPCircuitclass

AbstractPCircuit()

Abstract base class for probabilistic circuits.

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, '']) -> list[PyTree[Array]]

Default per-gate parameters, aligned with the circuit's gates.

Splits key and calls each gate's init_params.

Arguments:

key: PRNG key.

Returns:

A list of per-gate parameter pytrees, one per gate in circuit order.

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.

Concrete classes

DiscretePCircuitclass

DiscretePCircuit(gates: list[AbstractDiscreteGate], reps: int = 1)

Probabilistic circuit for discrete gates.

constructor

Initialize a probabilistic circuit.

Arguments:

gates: the list of gates to be applied in the circuit
reps: the number of times the circuit is applied to the initial input

gatesattribute

gates: list[AbstractDiscreteGate]

num_pditsattribute

num_pdits: int

dimsattribute

dims: tuple[int, ...]

repsattribute

reps: int

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.

init_paramsmethod

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

Default per-gate parameters, aligned with the circuit's gates.

Splits key and calls each gate's init_params.

Arguments:

key: PRNG key.

Returns:

A list of per-gate parameter pytrees, one per gate in circuit order.

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.

HybridStateclass

HybridState(*args, **kwargs)

Discrete and continuous registers of a hybrid circuit.

discreteattribute

discrete: <class 'jaxInt[Array, '... num_discrete_sites']'>

continuousattribute

continuous: <class 'jaxFloat[Array, '... continuous_dim']'>

HybridPCircuitclass

HybridPCircuit(gates: list[AbstractDiscreteGate | AbstractHybridGate], reps: int = 1)

Circuit supporting discrete, continuous, and hybrid gates.

Stores both discrete and continuous dimensions separately.

constructor

Initialize a hybrid circuit.

Arguments:

gates: List of gates (discrete, continuous, or hybrid).
reps: Number of times the circuit is repeated.

gatesattribute

gates: list[AbstractDiscreteGate | AbstractHybridGate]

discrete_dimsattribute

discrete_dims: tuple[int, ...]

continuous_dimsattribute

continuous_dims: tuple[int, ...]

repsattribute

reps: int

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.

init_paramsmethod

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

Default per-gate parameters, aligned with the circuit's gates.

Splits key and calls each gate's init_params.

Arguments:

key: PRNG key.

Returns:

A list of per-gate parameter pytrees, one per gate in circuit order.

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.

num_discrete_sitesproperty

num_discrete_sites

Number of discrete sites in the circuit.

num_continuous_sitesproperty

num_continuous_sites

Number of continuous sites in the circuit.

continuous_state_dimproperty

continuous_state_dim

Total dimension of continuous state = sum(continuous_dims).

continuous_offsetsproperty

continuous_offsets

Per continuous site (site, start, stop) slice into the register.