Parametrised stochastic circuits
Programs that transform probability distributions instead of fixed values. Build circuits from gates on three primitives, then sample them or propagate their moments exactly.
Programs that transform probability distributions instead of fixed values. Build circuits from gates on three primitives, then sample them or propagate their moments exactly.
Install with pip install extro-torx (Python 3.11+), then build and sample a circuit:
import jax
import jax.numpy as jnp
from torx.psc import DiscretePCircuit, BranchingSimulator, PSWAP
# PSWAP swaps two pbits with probability sigma(theta); here p(swap) = 0.3.
# Gates are structure-only; parameters are a separate list of logit leaves.
circuit = DiscretePCircuit([PSWAP([0, 1])])
thetas = [jnp.array([jnp.log(0.3 / 0.7)])]
sim = BranchingSimulator(num_samples=20_000)
compiled = sim.build_circuit(circuit, thetas)
# Start in |10) and sample the two output pbits.
state10 = jnp.array([1, 0], dtype=jnp.int32)
state01 = jnp.array([0, 1], dtype=jnp.int32)
samples = sim.sample(compiled, state10, jax.random.key(0))
stay = jnp.mean(jnp.all(samples == state10, axis=1))
swap = jnp.mean(jnp.all(samples == state01, axis=1))
print(f"stay |10): {float(stay):.3f}, swap |01): {float(swap):.3f}")stay |10): 0.699, swap |01): 0.30116 runnable notebooks, starting from a one-gate circuit and working up to trained networks and directed factor graphs.