Gates

Gates are stochastic kernels acting on pbits, pdits, and pmodes.

Abstract base classes

AbstractPGateclass

AbstractPGate()

Base class for all probabilistic gates.

For visualization, there are additional properties:

_control_indices: Indices of control sites for drawing.
() empty tuple: no controls, draw as spanning box (default)
(0,) tuple of ints: specific site indices are controls
"all_but_last": all sites except the last are controls
_draw_label: Optional short label for visualization (defaults to class name).

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, '']) -> ~_ThetaType

Return a freshly-initialised parameter pytree (theta).

Arguments:

key: PRNG key.

Returns:

An initial theta pytree for this gate.

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.

AbstractDiscreteGateclass

AbstractDiscreteGate()

A discrete probabilistic gate, a finite-state matrix factor.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Get the matrix representation of the gate.

Arguments:

theta: the gate's parameters.

Returns:

The column-stochastic transition matrix P[out, in] of the gate.

AbstractGeneratorGateclass

AbstractGeneratorGate()

Discrete gate defined by a continuous-time Markov generator $Q$.

Subclasses provide a rate matrix $Q$ via get_generator and a time step via dt. The transition matrix is then

$$P = \exp(Q \cdot \Delta t).$$

dtmethod

dt(theta: ~_ThetaType) -> Float[Array, '']

The continuous-time step $\Delta t$, extracted from theta.

get_generatormethod

get_generator(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Return the continuous-time Markov generator $Q$.

$Q$ must be a rate matrix, which means off-diagonal entries are non-negative and column sums equal to zero.

Arguments:

theta: the gate's parameters.

Returns:

The Markov generator $Q$.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Column-stochastic transition matrix $P = \exp(Q \cdot \Delta t)$.

See AbstractDiscreteGate.get_matrix for documentation.

AbstractHybridGateclass

AbstractHybridGate()

Base class for hybrid gates.

Gates specify which discrete and continuous sites they act on via the sites dict. As a Factor, the gate's sample takes the relevant substate as its inputs (a dict with "discrete" / "continuous" values at the gate's sites) and returns the new continuous substate.

Note

torx currently only supports discrete controlled continuous gates.

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 the new continuous substate.

Arguments:

key: JAX random key.
inputs: Substate dict with "discrete" / "continuous" keys

holding the values at the gate's sites.

params: the gate's theta.
info: optional runtime information.
site_info: unused.
return_aux: if True, return (output, None).

Returns:

New values for the continuous sites (same shape as inputs["continuous"]), or (output, None) when return_aux.

AbstractContinuousGateclass

AbstractContinuousGate()

Gate that only touches continuous sites.

Subclasses set sites to specify which continuous sites they act on.

constructor

Arguments:

sites: List of continuous site indices.
dims: Tuple of continuous dimensions for each site.

init_paramsmethod

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

Return a freshly-initialised parameter pytree (theta).

Arguments:

key: PRNG key.

Returns:

An initial theta pytree for this gate.

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

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 the new continuous substate.

Arguments:

key: JAX random key.
inputs: Substate dict with "discrete" / "continuous" keys

holding the values at the gate's sites.

params: the gate's theta.
info: optional runtime information.
site_info: unused.
return_aux: if True, return (output, None).

Returns:

New values for the continuous sites (same shape as inputs["continuous"]), or (output, None) when return_aux.

AbstractAffineGaussianGateclass

AbstractAffineGaussianGate()

Continuous gate with an exact affine-Gaussian channel.

Subclasses expose their local continuous-state transition as

$$ x' = A x + b + \epsilon,\qquad \epsilon \sim \mathcal{N}(0, \operatorname{diag}(\exp(\text{log\_var}))), $$

by implementing affine_parameters. The default sample draws from this channel, but exposing $(A, b, \log\text{var})$ also lets simulators reason about the exact Gaussian transition rather than only drawing samples.

constructor

Arguments:

sites: List of continuous site indices.
dims: Tuple of continuous dimensions for each site.

affine_parametersmethod

affine_parameters(
    theta: ~_ThetaType,
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return (A, b, log_var) for the local affine-Gaussian channel.

Arguments:

theta: the gate's parameters.

Returns:

A tuple (A, b, log_var) describing x' = A x + b + noise with noise ~ Normal(0, diag(exp(log_var))).

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 the affine-Gaussian channel exposed by affine_parameters.

AbstractControlledContinuousGateclass

AbstractControlledContinuousGate()

Gate where discrete sites control the continuous transformation.

Note

Currently we only support control by a single discrete site (pbit/pdit), so dims only specifies continuous dimensions. Multi-site discrete control would require tracking discrete dims as well.

constructor

Arguments:

sites: Tuple of (discrete_site, continuous_sites).
dims: Tuple of continuous dimensions.

init_paramsmethod

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

Return a freshly-initialised parameter pytree (theta).

Arguments:

key: PRNG key.

Returns:

An initial theta pytree for this gate.

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

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 the new continuous substate.

Arguments:

key: JAX random key.
inputs: Substate dict with "discrete" / "continuous" keys

holding the values at the gate's sites.

params: the gate's theta.
info: optional runtime information.
site_info: unused.
return_aux: if True, return (output, None).

Returns:

New values for the continuous sites (same shape as inputs["continuous"]), or (output, None) when return_aux.

AbstractKBranchGateclass

AbstractKBranchGate()

Gates with K branches selected via softmax over K-1 parameters.

The probability of branch k is softmax([0, theta])[k], where theta has K-1 learnable parameters for K branches. For K = 2, this reduces to the sigmoid parameterization: softmax([0, theta]) = [1 - sigmoid(theta), sigmoid(theta)].

Subclasses must implement:

num_branches: The number of branches K (>= 2)

branches: A list of K lookup tables, one per branch

num_branchesproperty

num_branches

Return the number of branches K (must be >= 2).

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

branchesproperty

branches

Get the K lookup tables for all branches as a stacked array.

Each lookup table (along axis 0) maps computational basis state index to output state values for each site. For example, for a PCNOT gate the operation branch is a CNOT gate. A CNOT gate's action on the computational basis states is 00 -> 00, 01 -> 01, 10 -> 11, and 11 -> 10, so the lookup table is [[0, 0], [0, 1], [1, 1], [1, 0]].

The identity branch (often branches[0]) maps each state to itself. For the PCNOT example, the identity maps 00 -> 00, 01 -> 01, 10 -> 10, 11 -> 11, so the lookup table is [[0, 0], [0, 1], [1, 0], [1, 1]].

For binary gates, basis_size = 2**num_sites. For pdit gates, basis_size = prod(dims).

Returns:

A stacked array of shape (K, basis_size, num_sites) containing all branch lookup tables.

AbstractSingleBinaryPGateclass

AbstractSingleBinaryPGate()

Abstract base class for single-site binary gates (dimension 2).

constructor

Arguments:

sites: The index of the site that this gate acts on.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Get the matrix representation of the gate.

Arguments:

theta: the gate's parameters.

Returns:

The column-stochastic transition matrix P[out, in] of the gate.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

num_branchesproperty

num_branches

Return the number of branches K (must be >= 2).

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

branchesproperty

branches

Get the K lookup tables for all branches as a stacked array.

For binary gates, basis_size = 2**num_sites. For pdit gates, basis_size = prod(dims).

Returns:

A stacked array of shape (K, basis_size, num_sites) containing all branch lookup tables.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

AbstractMultiBinaryPGateclass

AbstractMultiBinaryPGate()

Abstract base class for multi-site binary gates (dims = 2).

constructor

Arguments:

sites: A list of site indices that this gate acts on.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Get the matrix representation of the gate.

Arguments:

theta: the gate's parameters.

Returns:

The column-stochastic transition matrix P[out, in] of the gate.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

num_branchesproperty

num_branches

Return the number of branches K (must be >= 2).

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

branchesproperty

branches

Get the K lookup tables for all branches as a stacked array.

For binary gates, basis_size = 2**num_sites. For pdit gates, basis_size = prod(dims).

Returns:

A stacked array of shape (K, basis_size, num_sites) containing all branch lookup tables.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

AbstractSinglePditGateclass

AbstractSinglePditGate()

Abstract base class for single-site pdit gates (arbitrary dimension k).

constructor

Arguments:

sites: The index of the site that this gate acts on.
dims: Tuple containing the dimension of the pdit.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Get the matrix representation of the gate.

Arguments:

theta: the gate's parameters.

Returns:

The column-stochastic transition matrix P[out, in] of the gate.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

num_branchesproperty

num_branches

Return the number of branches K (must be >= 2).

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

branchesproperty

branches

Get the K lookup tables for all branches as a stacked array.

For binary gates, basis_size = 2**num_sites. For pdit gates, basis_size = prod(dims).

Returns:

A stacked array of shape (K, basis_size, num_sites) containing all branch lookup tables.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

AbstractMultiPditGateclass

AbstractMultiPditGate()

Abstract base class for multi-site pdit gates (same dimension k per site).

constructor

Arguments:

sites: A list of site indices that this gate acts on.
dims: Tuple containing the dimensions of each pdit.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Get the matrix representation of the gate.

Arguments:

theta: the gate's parameters.

Returns:

The column-stochastic transition matrix P[out, in] of the gate.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

num_branchesproperty

num_branches

Return the number of branches K (must be >= 2).

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

branchesproperty

branches

Get the K lookup tables for all branches as a stacked array.

For binary gates, basis_size = 2**num_sites. For pdit gates, basis_size = prod(dims).

Returns:

A stacked array of shape (K, basis_size, num_sites) containing all branch lookup tables.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

Concrete classes

HybridSitesclass

HybridSites(discrete: list[int], continuous: list[int])

Site mapping for hybrid gates.

Keys:

discrete: discrete site indices read as controls.
continuous: continuous site indices updated by the gate.

Continuous-only gates also accept a list of continuous site indices; hybrid gate constructors accept a (discrete, continuous) pair.

discreteattribute

discrete: list[int]

continuousattribute

continuous: list[int]

PNOTclass

PNOT(sites: int)

A probabilistic NOT gate.

This gate flips the value of a given pbit with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing.

The transition matrix is:

$$\begin{pmatrix} 1-p & p \\ p & 1-p \end{pmatrix}$$

constructor

See AbstractSingleBinaryPGate.__init__.

sitesattribute

sites: int

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 0 -> 0 and 1 -> 1. branches[1] is bit flip: 0 -> 1 and 1 -> 0.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '2 2']

See AbstractDiscreteGate.get_matrix for documentation.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PCNOTclass

PCNOT(sites: list[int])

A probabilistic CNOT gate.

This gate performs a controlled-NOT operation on two given pbits with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing. A controlled-NOT operation flips the value of the second pbit if the first pbit has the value 1, and does nothing if the first pbit has the value 0.

Transition matrix:

$$\begin{pmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 1-p & p \\ 0 & 0 & p & 1-p \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 00 -> 00, 01 -> 01, 10 -> 10, 11 -> 11. branches[1] is CNOT: 00->00, 01->01, 10->11, 11->10.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '4 4']

See AbstractDiscreteGate.get_matrix for documentation.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PSWAPclass

PSWAP(sites: list[int])

A probabilistic SWAP gate.

This gate swaps the values of two given pbits with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing.

Transition matrix:

$$\begin{pmatrix} 1 & 0 & 0 & 0 \\ 0 & 1-p & p & 0 \\ 0 & p & 1-p & 0 \\ 0 & 0 & 0 & 1 \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 00 -> 00, 01 -> 01, 10 -> 10, 11 -> 11. branches[1] is swap: 00->00, 01->10, 10->01, 11->11.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '4 4']

See AbstractDiscreteGate.get_matrix for documentation.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PJUMPclass

PJUMP(sites: list[int])

A probabilistic JUMP gate.

This gate performs a "jump" operation on two given pbits with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing. A jump operation moves probability from $|10)$ to $|01)$ if the first pbit has the value 1.

Transition matrix:

$$\begin{pmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & p & 0 \\ 0 & 0 & 1-p & 0 \\ 0 & 0 & 0 & 1 \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 00 -> 00, 01 -> 01, 10 -> 10, 11 -> 11. branches[1] is JUMP: 00->00, 01->01, 10->01, 11->11.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '4 4']

See AbstractDiscreteGate.get_matrix for documentation.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PMultiCNOTclass

PMultiCNOT(sites: list[int])

A probabilistic multi-controlled-NOT gate.

This gate performs a multi-controlled-NOT operation on any number of given pbits with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing. A multi-controlled-NOT operation flips the value of the last pbit if all the other pbits have the value 1, and does nothing otherwise.

The transition matrix is identity except for the last two rows/columns which swap with probability $p$.

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity. branches[1] is multi-controlled-NOT: flips last bit when all others are 1.

output_specproperty

output_spec

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, 'd d']

See AbstractDiscreteGate.get_matrix for documentation.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PDEMUXclass

PDEMUX(sites: list[int])

A probabilistic demultiplexor.

This gate has one input pbit and one "work" pbit. It either (i) copies the input to the work pbit and sets the input to 0, or (ii) does nothing to the input and sets the work pbit to 0. The probability of the former occurring is $p = \sigma(\theta)$ where $\theta$ is the gate parameter.

Transition matrix:

$$\begin{pmatrix} 1 & 1 & 0 & 0 \\ 0 & 0 & p & p \\ 0 & 0 & 1-p & 1-p \\ 0 & 0 & 0 & 0 \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] resets second pbit: 00 -> 00, 01 -> 00, 10 -> 10, 11 -> 10. branches[1] copies and resets: 00->00, 01->00, 10->01, 11->01.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '4 4']

See AbstractDiscreteGate.get_matrix for documentation.

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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PORclass

POR(sites: list[int])

A probabilistic OR gate.

This gate performs a modified OR operation on two given pbits with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing. A modified OR operation sets the first pbit to the OR of the two input pbits, and sets the second pbit to 0.

Transition matrix:

$$\begin{pmatrix} 1 & 0 & 0 & 0 \\ 0 & 1-p & 0 & 0 \\ 0 & p & 1 & p \\ 0 & 0 & 0 & 1-p \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 00 -> 00, 01 -> 01, 10 -> 10, 11 -> 11. branches[1] is OR: 00->00, 01->10, 10->10, 11->10.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '4 4']

See AbstractDiscreteGate.get_matrix for documentation.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PResetclass

PReset(sites: int)

A probabilistic reset gate.

This gate sets the given pbit to 0 with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing.

The transition matrix is:

$$\begin{pmatrix} 1 & p \\ 0 & 1-p \end{pmatrix}$$

constructor

See AbstractSingleBinaryPGate.__init__.

sitesattribute

sites: int

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 0 -> 0 and 1 -> 1. branches[1] resets to 0: 0 -> 0 and 1 -> 0.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '2 2']

See AbstractDiscreteGate.get_matrix for documentation.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

PCopyclass

PCopy(sites: list[int])

A probabilistic copy gate.

This gate copies the first pbit to the second pbit with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter, and otherwise does nothing.

Transition matrix:

$$\begin{pmatrix} 1 & p & 0 & 0 \\ 0 & 1-p & 0 & 0 \\ 0 & 0 & 1-p & 0 \\ 0 & 0 & p & 1 \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: 00->00, 01->01, 10->10, 11->11. branches[1] copies first to second: 00->00, 01->00, 10->11, 11->11.

dimsproperty

dims

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '4 4']

See AbstractDiscreteGate.get_matrix for documentation.

PCSWAPclass

PCSWAP(sites: list[int])

A probabilistic controlled-SWAP gate (Fredkin gate).

This gate swaps the values of two target pbits only when the control pbit is in state $|1)$, with probability $p = \sigma(\theta)$ where $\theta$ is the gate parameter. If the control is $|0)$ or the gate is not applied (with probability $1-p$), nothing happens.

The first site is the control, the second and third sites are the targets to be swapped.

Transition matrix:

$$\begin{pmatrix} 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 1-p & p & 0 \\ 0 & 0 & 0 & 0 & 0 & p & 1-p & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 \end{pmatrix}$$

constructor

See AbstractMultiBinaryPGate.__init__.

sitesattribute

sites: list[int]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

dimsproperty

dims

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity. branches[1] is controlled-SWAP: swap targets if control=1.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, '8 8']

See AbstractDiscreteGate.get_matrix for documentation.

GaussianNoiseGateclass

GaussianNoiseGate(sites: HybridSites | list[int], dims: tuple[int, ...])

Additive Gaussian noise gate.

This gate adds independent Gaussian noise to the continuous state:

$x' = x + \mathcal{N}(0, \exp(\theta))$

The variance is parameterized as $\exp(\theta)$ to ensure positivity.

constructor

See AbstractAffineGaussianGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[int, ...]

init_paramsmethod

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

Initial theta (log-variance), zeros over the continuous dims.

affine_parametersmethod

affine_parameters(
    theta: Float[Array, 'd'],
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return the additive Gaussian channel as (A, b, log_var).

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

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 the affine-Gaussian channel exposed by affine_parameters.

AffineGaussianGateclass

AffineGaussianGate(sites: HybridSites | list[int], dims: tuple[int, ...])

Affine transformation with Gaussian noise.

This gate applies an affine transformation followed by Gaussian noise: $x' = Ax + b + \mathcal{N}(0, \sigma^2)$ where $\sigma^2 = \exp$(log_var)

The theta dict contains:

A: Transformation matrix of shape (d, d)
b: Bias vector of shape (d,)
log_var: Log-variance of shape (d,)

constructor

See AbstractAffineGaussianGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[int, ...]

init_paramsmethod

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

Initial identity affine: A = I, b = 0, log_var = 0.

affine_parametersmethod

affine_parameters(
    theta: dict[str, Array],
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return (A, b, log_var) for this affine-Gaussian channel.

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

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 the affine-Gaussian channel exposed by affine_parameters.

Displaceclass

Displace(sites: HybridSites | list[int], dims: tuple[int, ...])

Displace continuous sites by theta.

This deterministic affine-Gaussian channel applies

$$ x' = x + \theta. $$

constructor

See AbstractAffineGaussianGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[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, '']) -> Float[Array, 'd']

Initial displacement theta of zeros over the continuous dims.

affine_parametersmethod

affine_parameters(
    theta: Float[Array, 'd'],
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return the displacement channel as (A, b, log_var).

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

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 the affine-Gaussian channel exposed by affine_parameters.

Scaleclass

Scale(sites: HybridSites | list[int], dims: tuple[int, ...])

Scale continuous sites by $\exp(\theta)$.

This deterministic affine-Gaussian channel applies

$$ x'_i = \exp(\theta_i) x_i. $$

constructor

See AbstractAffineGaussianGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[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.

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.

init_paramsmethod

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

Initial log-scale theta of zeros (i.e. unit scale).

affine_parametersmethod

affine_parameters(
    theta: Float[Array, 'd'],
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return the scaling channel as (A, b, log_var).

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

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 the affine-Gaussian channel exposed by affine_parameters.

Mixclass

Mix(sites: HybridSites | list[int], dims: tuple[int, ...])

Rotate exactly two scalar continuous sites by angle theta.

This deterministic affine-Gaussian channel applies

$$ \begin{bmatrix}x'_0\\x'_1\end{bmatrix} = \begin{bmatrix} \cos\theta & -\sin\theta\\ \sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix}x_0\\x_1\end{bmatrix}. $$

Expects two scalar continuous sites (dims = (1, 1)).

constructor

See AbstractAffineGaussianGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[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.

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

init_paramsmethod

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

Initial rotation angle theta of zero.

affine_parametersmethod

affine_parameters(
    theta: Float[Array, ''],
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return the rotation channel as (A, b, log_var).

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 the affine-Gaussian channel exposed by affine_parameters.

Diffuseclass

Diffuse(sites: HybridSites | list[int], dims: tuple[int, ...])

Apply Brownian diffusion with log-variance theta.

This additive affine-Gaussian channel applies

$$ x' = x + \epsilon,\qquad \epsilon \sim \mathcal{N}(0, \exp(\theta)). $$

For a diffusion coefficient $D$ over time $t$, set theta to $\log(2 D t)$. Functionally this matches GaussianNoiseGate; it is provided as a named Brownian-diffusion channel.

constructor

See AbstractAffineGaussianGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[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.

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.

discrete_dimsproperty

discrete_dims

Dimensions of discrete control sites, aligned with sites["discrete"].

input_portsproperty

input_ports

output_specproperty

output_spec

init_paramsmethod

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

Initial log-variance theta of zeros over the continuous dims.

affine_parametersmethod

affine_parameters(
    theta: Float[Array, 'd'],
) -> tuple[Float[Array, 'local_dim local_dim'], Float[Array, 'local_dim'], Float[Array, 'local_dim']]

Return the Brownian diffusion channel as (A, b, log_var).

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 the affine-Gaussian channel exposed by affine_parameters.

MixtureGaussianGateclass

MixtureGaussianGate(
    sites: HybridSites | tuple[int | list[int], int | list[int]],
    dims: tuple[int, ...],
    num_components: int,
)

Mixture Gaussian gate controlled by a discrete site.

This gate samples from one of K Gaussian components, where the component is selected by the discrete control site. Given discrete state $k$: $x' = x + \mu_k + \mathcal{N}(0, \sigma_k^2)$ where $\sigma_k^2 = \exp$(log_vars[k])

The discrete site should take values in $\{0, 1, ..., K-1\}$ where K is the number of mixture components.

Note

dims only specifies continuous dimensions since we only support control by a single discrete site for now.

The theta dict contains:

means: Component means of shape (K, d)
log_vars: Component log-variances of shape (K, d)

constructor

See AbstractControlledContinuousGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[int, ...]

num_componentsattribute

num_components: 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.

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.

input_portsproperty

input_ports

output_specproperty

output_spec

discrete_dimsproperty

discrete_dims

init_paramsmethod

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

Initial means and log_vars of zeros, shape (K, d).

samplemethod

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

Sample the new continuous substate.

Arguments:

key: JAX random key.
inputs: Substate dict with "discrete" / "continuous" keys

holding the values at the gate's sites.

params: the gate's theta.
info: optional runtime information.
site_info: unused.
return_aux: if True, return (output, None).

Returns:

New values for the continuous sites (same shape as inputs["continuous"]), or (output, None) when return_aux.

JumpDiffusionGateclass

JumpDiffusionGate(
    sites: HybridSites | tuple[int | list[int], int | list[int]],
    dims: tuple[int, ...],
)

Jump-diffusion gate controlled by a discrete firing site.

This gate applies continuous diffusion, and conditionally applies a jump based on the discrete control site. Given discrete state $j \in \{0, 1\}$:

$$ x' = x + \mathcal{N}(0, \sigma_d^2) + j \cdot (\mu_j + \mathcal{N}(0, \sigma_j^2)) $$

where $\sigma_d^2 = \exp$(diff_log_var) and $\sigma_j^2 = \exp$(jump_log_var).

When $j = 0$, only diffusion is applied. When $j = 1$, both diffusion and the jump are applied.

Note

dims only specifies continuous dimensions since we only support control by a single discrete site for now.

The theta dict contains:

diff_log_var: Diffusion log-variance of shape (d,)
jump_mean: Jump mean of shape (d,)
jump_log_var: Jump log-variance of shape (d,)

constructor

See AbstractControlledContinuousGate.__init__.

sitesattribute

sites: HybridSites

dimsattribute

dims: tuple[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.

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.

input_portsproperty

input_ports

output_specproperty

output_spec

discrete_dimsproperty

discrete_dims

init_paramsmethod

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

Initial diff_log_var/jump_mean/jump_log_var of zeros.

samplemethod

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

Sample the new continuous substate.

Arguments:

key: JAX random key.
inputs: Substate dict with "discrete" / "continuous" keys

holding the values at the gate's sites.

params: the gate's theta.
info: optional runtime information.
site_info: unused.
return_aux: if True, return (output, None).

Returns:

New values for the continuous sites (same shape as inputs["continuous"]), or (output, None) when return_aux.

PISINGclass

PISING(sites: list[int])

A probabilistic Ising interaction gate.

Unlike the Boolean gates (PNOT, PCNOT, PSWAP, ...) which interpolate between identity and a logic operation, PISING is an energy-based gate: it drives two pbits toward the thermal equilibrium of a pairwise Ising interaction. Applying it repeatedly thermalizes the joint distribution toward the Boltzmann distribution $\pi_i \propto e^{-\beta E_i}$.

Simulator support

PISING is a generator-defined dense-matrix gate. Use it with [StateVectorSimulator][torx.psc.StateVectorSimulator]; it is not compatible with [BranchingSimulator][torx.psc.BranchingSimulator], which currently requires branch lookup-table gates.

The Ising energy for two pbits $\sigma_1, \sigma_2 \in \{0, 1\}$ is:

$$E(\sigma_1, \sigma_2) = -J \, s_1 s_2 - h_1 s_1 - h_2 s_2$$

where $s_i = 2\sigma_i - 1 \in \{-1, +1\}$.

The transition matrix $P = \exp(Q \cdot \Delta t)$ is the matrix exponential of a Markov generator $Q$ whose off-diagonal entries are Glauber single-spin-flip rates:

$$Q_{ji} = \frac{1}{1 + e^{\beta \, \Delta E_{ij}}}$$

Only single-spin-flip transitions are nonzero.

Parameters:

theta[0], $J$: coupling strength

($J > 0$ ferromagnetic, $J < 0$ antiferromagnetic)

theta[1], $h_1$: external field on first pbit
theta[2], $h_2$: external field on second pbit
theta[3], $\beta$: inverse temperature

($\beta \to 0$: uniform, $\beta \to \infty$: ground state)

theta[4], $\Delta t$: continuous-time step

($0$: identity, larger: closer to equilibrium)

constructor

See AbstractDiscreteGate.__init__.

sitesattribute

sites: list[int]

dimsproperty

dims

init_paramsmethod

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

Initial theta = [J, h1, h2, beta, dt] (zeros).

See AbstractPGate.init_params.

Jmethod

J(theta: Float[Array, '5']) -> Float[Array, '']

h1method

h1(theta: Float[Array, '5']) -> Float[Array, '']

h2method

h2(theta: Float[Array, '5']) -> Float[Array, '']

betamethod

beta(theta: Float[Array, '5']) -> Float[Array, '']

dtmethod

dt(theta: Float[Array, '5']) -> Float[Array, '']

The continuous-time step $\Delta t$, extracted from theta.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

get_generatormethod

get_generator(theta: Float[Array, '5']) -> Float[Array, '4 4']

Build the 4x4 Glauber-dynamics Markov generator $Q$.

$Q$ is a rate matrix with zero column sums and non-negative off-diagonal entries. Only single-spin-flip transitions are nonzero: $|00) \leftrightarrow |01)$, $|00) \leftrightarrow |10)$, $|01) \leftrightarrow |11)$, $|10) \leftrightarrow |11)$.

The generator is independent of $\Delta t$.

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

get_matrixmethod

get_matrix(theta: ~_ThetaType) -> Float[Array, 'dim1 dim2']

Column-stochastic transition matrix $P = \exp(Q \cdot \Delta t)$.

See AbstractDiscreteGate.get_matrix for documentation.

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.

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_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.

PditShiftclass

PditShift(sites: int, dims: Any)

A probabilistic cyclic shift gate for a single k-dimensional pdit.

This gate shifts the state cyclically: i -> (i+1) mod k with probability $p = \sigma(\theta)$, and does nothing with probability $1-p$.

The transition matrix for dimension k is: $(1 - p) I + p C$ where $C$ is the cyclic permutation matrix.

constructor

See AbstractSinglePditGate.__init__.

sitesattribute

sites: int

dimsattribute

dims: tuple[int, ...]

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: i -> i. branches[1] is cyclic shift: i -> (i+1) mod k.

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, 'd d']

See AbstractDiscreteGate.get_matrix for documentation.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

PditSWAPclass

PditSWAP(sites: list[int], dims: Any)

A probabilistic SWAP gate for two k-dimensional pdits.

This gate swaps the values of two pdits with probability $p = \sigma(\theta)$, and does nothing with probability $1-p$. Both pdits must have the same dimension.

For dimension k, the state space is $k^2$ and the SWAP permutes basis states $|i,j) \to |j,i)$.

constructor

See AbstractMultiPditGate.__init__.

sitesattribute

sites: list[int]

dimsattribute

dims: tuple[int, ...]

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: |i,j) -> |i,j). branches[1] is SWAP: |i,j) -> |j,i).

get_matrixmethod

get_matrix(theta: Float[Array, '1']) -> Float[Array, 'd2 d2']

See AbstractDiscreteGate.get_matrix for documentation.

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.

PditCycleclass

PditCycle(sites: int, dims: Any)

A probabilistic 3-branch cyclic gate for a single d-dimensional pdit.

This gate implements a random walk on a ring: with some probability stay in place, shift forward, or shift backward. Useful for modeling diffusion or random walks on discrete rings.

This gate has 3 branches:

Branch 0: Identity (i -> i)
Branch 1: Forward cycle (i -> (i+1) mod d)
Branch 2: Backward cycle (i -> (i-1) mod d)

The probabilities are determined by softmax([0, theta[0], theta[1]]).

The transition matrix is:

$p_0 I + p_1 C_{forward} + p_2 C_{backward}$

where $C_{forward}$ is the forward cyclic permutation and $C_{backward}$ is the backward cyclic permutation.

constructor

See AbstractSinglePditGate.__init__.

sitesattribute

sites: int

dimsattribute

dims: tuple[int, ...]

input_statesproperty

input_states

output_statesproperty

output_states

input_portsproperty

input_ports

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']

Row-stochastic log P[in, out] = log(get_matrix(params).T).

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]]

Draw the gate's output configuration from its transition matrix.

num_branchesproperty

num_branches

branchesproperty

branches

branches[0] is identity: i -> i. branches[1] is forward cycle: i -> (i+1) mod d. branches[2] is backward cycle: i -> (i-1) mod d.

get_matrixmethod

get_matrix(theta: Float[Array, '2']) -> Float[Array, 'd d']

See AbstractDiscreteGate.get_matrix for documentation.

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.

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_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.

init_paramsmethod

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

Initial theta of shape (num_branches - 1,) (zeros).

See AbstractPGate.init_params.

probsmethod

probs(theta: Float[Array, 'num_branch_params']) -> Float[Array, 'K']

Branch probabilities for parameters theta.

probmethod

prob(theta: Float[Array, 'num_branch_params']) -> Float[Array, '']

Probability of applying the operation (branches[1]); K=2 only.