Stochastic Samplers¤
traceax uses a flexible approach to define how random samples are generated within
traceax.AbstractTraceEstimator instances. While this typically wraps a single
jax random call, the varied interfaces for each randomization procedure may differ,
which makes uniformly interfacing with it a bit annoying. As such, we provide a
simple abstract class definition, traceax.AbstractSampler using that subclasses
Equinox modules.
traceax.AbstractSampler
traceax.AbstractSampler
¤
Abstract base class for all samplers.
__call__(self, key: PRNGKeyArray, n: int, k: int) -> Num[Array, 'n k']
abstractmethod
¤
Sample random variates from the underlying distribution as an \(n \times k\) matrix.
Example
sampler = tr.RademacherSampler()
samples = sampler(key, n, k)
Each sampler accepts a dtype (i.e. float, complex, int) argument upon initialization,
with sensible default values. This makes it possible to sample from more general spaces (e.g.,
complex Normal test-vectors).
Example
sampler = tr.NormalSampler(complex)
samples = sampler(key, n, k)
Arguments:
key: a jax PRNG key used as the random key.n: the size of the leading dimension.k: the size of the trailing dimension.
Returns:
An Array of random samples.
traceax.NormalSampler(AbstractSampler)
¤
Standard normal distribution sampler.
Generates samples \(X_{ij} \sim N(0, 1)\) for \(i \in [n]\) and \(j \in [k]\).
Note
Supports float and complex-valued types.
__init__(dtype: DTypeLike = float)
¤
Arguments:
dtype: numeric representation for sampled test-vectors. Default isfloat.
traceax.SphereSampler(AbstractSampler)
¤
Sphere distribution sampler.
Generates samples \(X_1, \dotsc, X_n\) uniformly distributed on the surface of a \(k\) dimensional sphere (i.e. \(k-1\)-sphere) with radius \(\sqrt{n}\). Internally, this operates by sampling standard normal variates, and then rescaling such that each \(k\)-vector \(X_i\) has \(\lVert X_i \rVert = \sqrt{n}\).
Note
Supports float and complex-valued types.
__init__(dtype: DTypeLike = float)
¤
Arguments:
dtype: numeric representation for sampled test-vectors. Default isfloat.
traceax.RademacherSampler(AbstractSampler)
¤
Rademacher distribution sampler.
Generates samples \(X_{ij} \sim \mathcal{U}(-1, +1)\) for \(i \in [n]\) and \(j \in [k]\).
Note
Supports integer, float, and complex-valued types.
__init__(dtype: DTypeLike = int)
¤
Arguments:
dtype: numeric representation for sampled test-vectors. Default isint.