A Framework for Stochastic Differentiable Programming

Abstract

We introduce Parametrized Stochastic Circuits (PSCs), a gate-based intermediate representation for programmable stochastic dynamics in which typed local stochastic kernels with tunable parameters compose over explicit binary, categorical, and continuous wires, and \texttt{torx}, an open-source JAX framework for constructing, executing, and differentiating them. PSCs' data types and stochastic kernels are chosen to align closely with the native operations exposed by emerging probabilistic hardware. In this way, stochastic algorithms can be designed directly in terms of the operations the hardware executes natively, so that the energy advantage arising at this level is not lost on mappings that introduce substantial decomposition, communication, or control overhead. We demonstrate the framework on a variety of example applications such as random walks on graphs, discrete diffusion, stochastic graph networks, jump diffusion and Ising sampling. We also report a hardware experiment in which probabilistic bits on the X0 subthreshold CMOS test chip, hosted by the XTR-0 desktop platform, provide physical randomness for Metropolis-Hastings and importance-sampling estimators, yielding estimates consistent with a software pseudorandom baseline.