TORX

Simulators

A simulator compiles a circuit and reads it back as samples, moments, or a density.

Abstract base classes

AbstractSamplerclass
AbstractSampler()

Produces random draws used by sample-based simulators.

bernoullimethod
bernoulli(
    key: Key[Array, ''],
    p: Float[Array, '...'],
    shape: tuple[int, ...] | None = None,
) -> Int[Array, '...']
categoricalmethod
categorical(
    key: Key[Array, ''],
    logits: Float[Array, '... K'],
    axis: int = -1,
    shape: tuple[int, ...] | None = None,
) -> Int[Array, '...']
normalmethod
normal(
    key: Key[Array, ''],
    shape: tuple[int, ...] = (),
    dtype: Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType] = <class 'float'>,
) -> Float[Array, '...']
AbstractCompiledPCircuitclass
AbstractCompiledPCircuit()

Abstract parent class for probabilistic circuits built for a specific backend.

from_pcircuitclassmethod
from_pcircuit(circuit: ~_CircuitType, thetas: PyTree[Array]) -> Self

Compile the given probabilistic circuit with the given parameters.

Arguments:

  • circuit: The probabilistic circuit to compile
  • thetas: The parameters to bake into the compiled circuit.

Returns:

The compiled circuit.

to_pcircuitmethod
to_pcircuit(structure: ~_CircuitType) -> ~_CircuitType

Return a circuit structure reconstructed by this compiled backend.

AbstractSimulatorclass
AbstractSimulator()

Abstract parent class for probabilistic circuit simulators.

expvalmethod
expval(*args, **kwargs)

Compute the expectation value of the given index after circuit execution.

expval_allmethod
expval_all(*args, **kwargs)

Compute the expectation value of all indices after circuit execution.

build_circuitmethod
build_circuit(circuit: ~_CircuitType, thetas: PyTree[Array]) -> ~_CompiledType

Build the circuit for this simulator with the given parameters.

Arguments:

  • circuit: The circuit to build
  • thetas: The parameters to bake into the compiled circuit.

Returns:

The built circuit.

Concrete classes

JaxPRNGSamplerclass
JaxPRNGSampler()

Default sampler drawing from jax.random with a JAX PRNG key.

bernoullimethod
bernoulli(
    key: Key[Array, ''],
    p: Float[Array, '...'],
    shape: tuple[int, ...] | None = None,
) -> Int[Array, '...']
categoricalmethod
categorical(
    key: Key[Array, ''],
    logits: Float[Array, '... K'],
    axis: int = -1,
    shape: tuple[int, ...] | None = None,
) -> Int[Array, '...']
normalmethod
normal(
    key: Key[Array, ''],
    shape: tuple[int, ...] = (),
    dtype: Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType] = <class 'float'>,
) -> Float[Array, '...']
GaussianMomentsclass
GaussianMoments(
    mean: Float[Array, 'continuous_dim'],
    covariance: Float[Array, 'continuous_dim continuous_dim'],
    site_offsets: tuple[tuple[int, int, int], ...],
    observed_sites: tuple[int, ...] = (),
)

Joint Gaussian moments over continuous sites.

mean and covariance follow the order given by sites. site_offsets maps each site to its (start, stop) slice into them.

meanattribute
mean: <class 'jaxFloat[Array, 'continuous_dim']'>
covarianceattribute
covariance: <class 'jaxFloat[Array, 'continuous_dim continuous_dim']'>
site_offsetsattribute
site_offsets: tuple[tuple[int, int, int], ...]
observed_sitesattribute
observed_sites: tuple[int, ...]
sitesproperty
sites

Sites covered by this state, in mean/covariance order.

site_indicesmethod
site_indices(site: int) -> Int[Array, 'd']

Return flat state-vector indices for one continuous site.

site_momentsmethod
site_moments(site: int) -> tuple[Float[Array, 'd'], Float[Array, 'd d']]

Return marginal mean and covariance for one continuous site.

CompiledAffineGaussianPCircuitclass
CompiledAffineGaussianPCircuit(
    gates: list[AbstractDiscreteGate | AbstractHybridGate],
    thetas: list[PyTree[Array]],
    site_offsets: tuple[tuple[int, int, int], ...],
    reps: int,
)

Compiled hybrid circuit for the affine Gaussian simulator.

gatesattribute
gates: list[AbstractDiscreteGate | AbstractHybridGate]
thetasattribute
thetas: list[jaxPyTree[Array]]
site_offsetsattribute
site_offsets: tuple[tuple[int, int, int], ...]
repsattribute
reps: int
from_pcircuitclassmethod
from_pcircuit(circuit: HybridPCircuit, thetas: list[PyTree[Array]]) -> Self

Compile circuit for exact affine Gaussian moment propagation.

Arguments:

  • circuit: The hybrid circuit to compile.
  • thetas: Per-gate parameters aligned with circuit.gates.

Returns:

The compiled circuit.

to_pcircuitmethod
to_pcircuit(structure: HybridPCircuit) -> HybridPCircuit

Return a HybridPCircuit with structure's gates and compiled reps.

AffineGaussianSimulatorclass
AffineGaussianSimulator()

Exact moment simulator for the affine Gaussian fragment of hybrid circuits.

Propagates the joint mean and covariance in closed form through each gate's affine Gaussian channel (A, b, log_var), exposed via [AbstractAffineGaussianGate][torx.psc.AbstractAffineGaussianGate], and conditions on observed sites via the Schur complement. Only supports Gaussian gates.

Dense reference implementation: O(D^3) per gate in the total continuous dimension D; intended for small affine-Gaussian circuits.

build_circuitmethod
build_circuit(
    circuit: HybridPCircuit,
    thetas: list[PyTree[Array]],
) -> CompiledAffineGaussianPCircuit

Compile circuit for exact affine Gaussian moment propagation.

propagatemethod
propagate(
    circuit: CompiledAffineGaussianPCircuit,
    initial_continuous: Float[Array, 'continuous_dim'],
) -> GaussianMoments

Propagate joint Gaussian moments through an affine-Gaussian circuit.

Arguments:

  • circuit: The compiled affine Gaussian circuit to execute.
  • initial_continuous: Flat initial continuous state.

Returns:

GaussianMoments over every continuous site.

conditionmethod
condition(
    circuit: CompiledAffineGaussianPCircuit,
    observations: Mapping[int, Array] | None = None,
    *,
    initial_continuous: Float[Array, 'continuous_dim'],
    query_sites: Sequence[int] | None = None,
    jitter: Float[Array, ''] | float = 0.0,
) -> GaussianMoments

Condition final Gaussian moments on continuous-site observations.

The circuit is propagated to a joint Gaussian over the final continuous state, then the queried sites are conditioned on the observed ones via the Schur complement. The solve uses a Cholesky factorization of the observed covariance plus jitter on the diagonal. Observed/query site membership must be static Python-level dict keys or Sequence entries; only observation values and jitter may be traced, so JIT/vmap over which sites are observed is unsupported by design.

Arguments:

  • circuit: The compiled affine Gaussian circuit to execute.
  • observations: Mapping from each observed continuous site to its

    • value.

  • initial_continuous: Flat initial continuous state.
  • query_sites: Sites to return. If None, all unobserved sites are

    • queried.

  • jitter: Nonnegative diagonal regularizer for the observed

    • covariance.

Returns:

GaussianMoments over the queried sites, with observed_sites set.

expvalmethod
expval(
    circuit: CompiledAffineGaussianPCircuit,
    initial_continuous: Float[Array, 'continuous_dim'],
    site: int = 0,
) -> Float[Array, 'd']

Return the marginal mean of one continuous site.

Arguments:

  • circuit: The compiled affine Gaussian circuit to execute.
  • initial_continuous: Flat initial continuous state.
  • site: The continuous site whose marginal mean to return.

Returns:

The marginal mean of site.

expval_allmethod
expval_all(
    circuit: CompiledAffineGaussianPCircuit,
    initial_continuous: Float[Array, 'continuous_dim'],
) -> Float[Array, 'continuous_dim']

Return the joint mean over all continuous sites.

HybridStateclass
HybridState(*args, **kwargs)

State representation for hybrid circuits.

discreteattribute
discrete: <class 'jaxInt[Array, '... num_discrete_sites']'>
continuousattribute
continuous: <class 'jaxFloat[Array, '... continuous_dim']'>
CompiledHybridPCircuitclass
CompiledHybridPCircuit(
    gates: list[AbstractDiscreteGate | AbstractHybridGate],
    thetas: list[PyTree[Array]],
    discrete_dims: tuple[int, ...],
    continuous_dims: tuple[int, ...],
    reps: int,
    sampler: AbstractSampler,
)

Compiled hybrid circuit for the hybrid sample simulator.

gatesattribute
gates: list[AbstractDiscreteGate | AbstractHybridGate]
thetasattribute
thetas: list[jaxPyTree[Array]]
discrete_dimsattribute
discrete_dims: tuple[int, ...]
continuous_dimsattribute
continuous_dims: tuple[int, ...]
repsattribute
reps: int
samplerattribute
sampler: AbstractSampler
from_pcircuitclassmethod
from_pcircuit(
    circuit: HybridPCircuit,
    thetas: list[PyTree[Array]],
    sampler: AbstractSampler | None = None,
) -> Self

Compile a hybrid circuit with the given parameters.

Arguments:

  • circuit: The hybrid circuit to compile.
  • thetas: Per-gate parameters aligned with circuit.gates.
  • sampler: Optional entropy source for both discrete branch

    • decisions and continuous Gaussian noise.

Returns:

The compiled circuit.

to_pcircuitmethod
to_pcircuit(structure: HybridPCircuit) -> HybridPCircuit

Return a HybridPCircuit with this compiled circuit's gates and reps.

HybridSampleSimulatorclass
HybridSampleSimulator(num_samples: int = 1, sampler: AbstractSampler | None = None)

Sample-based simulator for hybrid circuits.

Handles both discrete gates (via branch sampling) and continuous/hybrid gates (via their sample() method).

Compilation time

This simulator unrolls the full loop over gates during tracing, so compilation time scales poorly with the number of gates.

constructor

Initialize the hybrid sample simulator.

Arguments:

  • num_samples: number of samples used to estimate expectation values
  • sampler: source of randomness for both discrete branch decisions

    • and continuous Gaussian noise.

num_samplesattribute
num_samples: int
samplerattribute
sampler: AbstractSampler
build_circuitmethod
build_circuit(
    circuit: HybridPCircuit,
    thetas: list[PyTree[Array]],
) -> CompiledHybridPCircuit

Compile circuit with this simulator's branch sampler attached.

samplemethod
sample(
    circuit: CompiledHybridPCircuit,
    initial_state: HybridState,
    key: Key[Array, ''],
) -> HybridState

Run circuit and return samples of final state.

Arguments:

  • circuit: The compiled hybrid circuit to execute.
  • initial_state: Initial state with "discrete" and "continuous" arrays.
  • key: JAX random key.

Returns:

Dict with "discrete" and "continuous" final state arrays.

expvalmethod
expval(
    circuit: CompiledHybridPCircuit,
    initial_state: HybridState,
    site: int,
    key: Key[Array, ''],
    site_type: str,
) -> PyTree[Array]

Estimate expectation value of a site.

Arguments:

  • circuit: The hybrid circuit to execute.
  • initial_state: Initial state with "discrete" and "continuous" arrays.
  • site: Index of the site.
  • key: JAX random key.
  • site_type: Either "discrete" or "continuous".

Returns:

The scalar expectation value of the discrete site, or the marginal mean vector of the continuous site.

expval_allmethod
expval_all(
    circuit: CompiledHybridPCircuit,
    initial_state: HybridState,
    key: Key[Array, ''],
) -> HybridState

Estimate expectation values of all sites.

Arguments:

  • circuit: The hybrid circuit to execute.
  • initial_state: Initial state with "discrete" and "continuous" arrays.
  • key: JAX random key.

Returns:

Dict with "discrete" and "continuous" expectation value arrays.

CompiledSamplePCircuitclass
CompiledSamplePCircuit(
    num_pdits: int,
    reps: int,
    max_branches: int,
    branch_ops: Int[Array, 'num_gates max_branches max_basis max_l'],
    num_branches: Int[Array, 'num_gates'],
    sites: Int[Array, 'num_gates l'],
    dims: Int[Array, 'num_gates l'],
    basis_sizes: Int[Array, 'num_gates'],
    thetas: Float[Array, 'num_gates max_branches_minus_1'],
    sampler: AbstractSampler,
)

Compiled probabilistic circuit class for the sample-based simulator.

num_pditsattribute
num_pdits: int
repsattribute
reps: int
max_branchesattribute
max_branches: int
branch_opsattribute
branch_ops: <class 'jaxInt[Array, 'num_gates max_branches max_basis max_l']'>
num_branchesattribute
num_branches: <class 'jaxInt[Array, 'num_gates']'>
sitesattribute
sites: <class 'jaxInt[Array, 'num_gates l']'>
dimsattribute
dims: <class 'jaxInt[Array, 'num_gates l']'>
basis_sizesattribute
basis_sizes: <class 'jaxInt[Array, 'num_gates']'>
thetasattribute
thetas: <class 'jaxFloat[Array, 'num_gates max_branches_minus_1']'>
samplerattribute
sampler: AbstractSampler
from_pcircuitclassmethod
from_pcircuit(
    circuit: DiscretePCircuit,
    thetas: list[Float[Array, '...']],
    sampler: AbstractSampler | None = None,
) -> Self

Compile the given probabilistic circuit with the given parameters.

This compiled form removes the list of probabilistic gates in favour of a more JIT-friendly representation. This representation uses the matrix forms of the branches of the probabilistic gates.

Arguments:

  • circuit: The probabilistic circuit to compile
  • thetas: Per-gate parameters aligned with circuit.gates; stacked and padded to (num_gates, max_branches - 1) with -inf.
  • sampler: Optional entropy source for branch decisions.

Returns:

The compiled circuit.

to_pcircuitmethod
to_pcircuit(structure: DiscretePCircuit) -> DiscretePCircuit

Return a DiscretePCircuit with structure's gates and compiled reps.

to_thetasmethod
to_thetas() -> list[Float[Array, '...']]

Recover the per-gate parameter list from the padded thetas.

Inverse of the stacking done in from_pcircuit: each gate's theta is the first K - 1 entries of its padded row.

Returns:

A list of per-gate thetas, aligned with the original gates.

SampleSimulatorclass
SampleSimulator(
    diff_method: Literal['param_shift_inf', 'param_shift_single', 'param_shift_filter'] = 'param_shift_inf',
    num_samples: int = 1,
    sampler: AbstractSampler | None = None,
)

A sample-based simulator for probabilistic circuits.

Instead of storing full state vectors, it samples from distributions.

Choosing a differentiation method

Three differentiation methods are available:

  • "param_shift_inf": Uses the parameter shift rule with deterministic

    • gates ($\theta \to \pm\infty$). Requires $2N$ circuit evaluations for $N$ parameters.

  • "param_shift_single": Uses the parameter shift rule with primal reuse.

    • Requires $N$ circuit evaluations.

  • "param_shift_filter": Estimates gradients from a single forward pass

    • by filtering samples based on which branch was taken at each gate. For each gate, samples are partitioned into those that applied the gate vs. those that did not, and the gradient is estimated from the difference in expectation values between these groups.

constructor

Initialize the sample simulator.

Arguments:

  • diff_method: method used for differentiating circuit parameters
  • num_samples: number of samples used to estimate expectation values
  • sampler: source of randomness for branch decisions.
diff_methodattribute
diff_method: Literal['param_shift_inf', 'param_shift_single', 'param_shift_filter']
num_samplesattribute
num_samples: int
samplerattribute
sampler: AbstractSampler
samplemethod
sample(
    circuit: CompiledSamplePCircuit,
    x: Int[Array, 'pbits'],
    key: Key[Array, ''],
) -> Int[Array, 'num_samples num_pbits']

Obtain samples from the final distribution of the probabilistic circuit.

Arguments:

  • circuit: The probabilistic circuit to execute
  • x: The initial computational basis state of the circuit
  • key: The random key to use to obtain samples

Returns:

An integer array with shape (num_samples, num_pbits) containing the computational basis state samples.

expvalmethod
expval(
    circuit: CompiledSamplePCircuit,
    x: Int[Array, 'pbits'],
    pbit: int,
    key: Key[Array, ''],
) -> Float[Array, '']

Estimate the expectation value of the given discrete site after circuit execution.

For a site with basis values $0, \ldots, d - 1$, this estimates $\mathbb{E}[x] = \sum_i i \, P(x = i)$.

Arguments:

  • circuit: The probabilistic circuit to execute
  • x: The initial computational basis state of the circuit
  • pbit: The index of the site to estimate the final expectation value of
  • key: The random key to use to obtain samples

Returns:

The expectation value of the given site after circuit execution.

expval_allmethod
expval_all(
    circuit: CompiledSamplePCircuit,
    x: Int[Array, 'pbits'],
    key: Key[Array, ''],
) -> Float[Array, 'num_pbits']

Estimate the expectation value of all discrete sites after circuit execution.

For each site with basis values $0, \ldots, d - 1$, this estimates $\mathbb{E}[x] = \sum_i i \, P(x = i)$.

Arguments:

  • circuit: The probabilistic circuit to execute
  • x: The initial computational basis state of the circuit
  • key: The random key to use to obtain samples

Returns:

An array containing the expectation values of all discrete sites.

build_circuitmethod
build_circuit(
    circuit: DiscretePCircuit,
    thetas: list[Float[Array, '...']],
) -> CompiledSamplePCircuit

Compile circuit with this simulator's branch sampler attached.

Arguments:

  • circuit: The probabilistic circuit to compile.
  • thetas: Per-gate parameters aligned with circuit.gates.

Returns:

The compiled circuit.

CompiledStateVectorPCircuitclass
CompiledStateVectorPCircuit(
    gates: list[AbstractDiscreteGate],
    thetas: list[Float[Array, '...']],
    num_pdits: int,
    dims: tuple[int, ...],
    reps: int,
)

Compiled probabilistic circuit class for the state vector simulator.

gatesattribute
gates: list[AbstractDiscreteGate]
thetasattribute
thetas: list[jaxFloat[Array, '...']]
num_pditsattribute
num_pdits: int
dimsattribute
dims: tuple[int, ...]
repsattribute
reps: int
from_pcircuitclassmethod
from_pcircuit(circuit: DiscretePCircuit, thetas: list[Float[Array, '...']]) -> Self

Compile the given probabilistic circuit with the given parameters.

Arguments:

  • circuit: The probabilistic circuit to compile
  • thetas: Per-gate parameters aligned with circuit.gates

Returns:

The compiled circuit.

to_pcircuitmethod
to_pcircuit(structure: DiscretePCircuit) -> DiscretePCircuit

Return a DiscretePCircuit with this compiled circuit's gates and reps.

StateVectorSimulatorclass
StateVectorSimulator()

Simulator for exact state vectors representing probability distributions.

build_circuitmethod
build_circuit(
    circuit: DiscretePCircuit,
    thetas: list[Float[Array, '...']],
) -> CompiledStateVectorPCircuit

Build the circuit for this simulator with the given parameters.

Arguments:

  • circuit: The circuit to build
  • thetas: The parameters to bake into the compiled circuit.

Returns:

The built circuit.

apply_gatestaticmethod
apply_gate(
    state: Float[Array, 'dimensions'],
    gate: AbstractDiscreteGate,
    theta: Float[Array, '...'],
    dims: tuple[int, ...],
) -> Float[Array, 'dimensions']

Apply gate to the state StateVector and return the resulting state.

Arguments:

  • state: The state to apply the gate to
  • gate: The gate to apply
  • theta: The gate's parameters
  • dims: The dimensions of all sites in the circuit

Returns:

The state after applying the gate.

densitymethod
density(
    circuit: CompiledStateVectorPCircuit,
    x: Float[Array, 'dimensions'],
) -> Float[Array, 'dimensions']

Compute the final distribution over computational basis states.

Arguments:

  • circuit: The probabilistic circuit to compute the distribution of
  • x: The initial state vector, also the initial distribution

Returns:

The final distribution of the circuit.

expvalmethod
expval(
    circuit: CompiledStateVectorPCircuit,
    x: Float[Array, 'dimensions'],
    pbit: int,
) -> Float[Array, '']

Compute the expectation value of the given discrete site after circuit execution.

For a site with basis values $0, \ldots, d - 1$, this returns $\sum_i i \, P(x = i)$. For binary sites, this is equivalent to the probability of measuring 1.

Arguments:

  • circuit: The probabilistic circuit to execute
  • x: The initial state vector, also the initial distribution
  • pbit: The index of the site to compute the final expectation value of

Returns:

The expectation value of the given site after circuit execution.

expval_allmethod
expval_all(
    circuit: CompiledStateVectorPCircuit,
    x: Float[Array, 'dimensions'],
) -> Float[Array, 'pbits']

Compute the expectation value of all discrete sites after circuit execution.

For each site with basis values $0, \ldots, d - 1$, this returns $\sum_i i \, P(x = i)$. For binary sites, this is equivalent to the probability of measuring 1.

Arguments:

  • circuit: The probabilistic circuit to execute
  • x: The initial state vector, also the initial distribution

Returns:

An array containing the expectation values of all discrete sites.