Available probability distribution in Aesara

[rlouf]

I gathered the following table which enumerates the distributions available in numpy or aesara.tensor.random.basic as well as whether they can be transpiled to Numba/JAX:

distribution	Numpy	Aesara	Aesara (Numba)	Aesara (JAX)
integers	x	x		#1331
random	x	x
choice	x	x	x	x
bytes	x	-	-	-
betabinomial		x		#1332
shuffle	x	#1367		x
permutation	x	x	x	x
permuted	x	#1354
bernoulli		x	x	x
beta	x	x	x	x
binomial	x	x	x	#1321
categorical		x	x	x
cauchy		x	x	x
chisquare	x	x	x	#1363
dirichlet	x	x	x	x
exponential	x	x	x	x
f	x	#1366
gamma	x	x	x	x
gengamma		x		#1333
geometric (geom)	x	x	x	x
gumbel	x	x	x	x
halfcauchy		x		x
halfnormal		x	x	x
hypergeometric (hypergeom)	x	x	x	#1324
invgamma		x		#1368
laplace	x	x	x	x
logistic	x	x	x	x
lognormal	x	x	x	x
logseries (logser)	x	#1349
multinomial	x	x	x	#1326
multivariate_hypergeometric	x	#1355
multivariate_normal	x	x		x
negative_binomial (nbinom)	x	x	x	#1327
noncentral_chisquare	x	#1356
noncentral_f	x	#1357
normal	x	x	x	x
pareto	x	x	x	x
poisson	x	x	x	x
power	x	#1352
rayleigh	x	x
randint		x	x	x
standard_cauchy	x	#1358
standard_exponential	x	#1358
standard_gamma	x	#1358
standard_normal	x	x		x
standard_t	x	#1358		x
t		x		x
triangular	x	x	x	#1328
uniform	x	x	x	x
vonmises	x	x	x	#1330
wald	x	x	x	x
weibull	x	x	x	x
zipf	x	#1350

• The list of distributions available in Numpy was retrieved from the documentation.
• Aesara transpiles RVs to Numba following the rules in this file.
• Aesara automatically transpiles RVs with name name to jax.random.name. List of available distributions in the JAX documentation.

Should we open issues for distributions that are missing from numpy / for which there is no Numba implementation?

A separate repo for probability distributions?

There are also many distributions that are implemented e.g. in scipy and that can be useful for bayesian inference (Wishart for instance). We probably don’t want to bloat Aesara with too many distributions, then I’d like to start a discussion about opening a separate repo (aerand?) with distributions that are frequently used in bayesian inference. This would have several advantages:

1. Keep the number distributions in Aesara to a minimum;
2. We could remove distribution-specific logic out of aeppl by moving the <name>_logprob and <name>_logcd funtions to this repo;
3. We could add extra information about random variables, such as their distribution’s support and constraints on their parameters (see MCX for an example of constraints implementation). This information can be used to remove distribution-specific logic for the transforms in aeppl, but also for static validation. For instance aeppl could return a warning when asked to build the logprob graph for a model where parameters have incompatible support. For instance a Normal prior on a standard deviation.

[brandonwillard]

A separate repo for probability distributions?

I've often thought that we should have a separate repo for SciPy Ops in general.

[ricardoV94]

We could implement some Numba / Jax distributions that don't have an immediate Op in a backend by using a more clever dispatch (or even a Aesara rewrite). For instance in JAX: Binomial(n, p) = sum(Bernoulli(p, size=n)) (they really don't have a binomial?)

[rlouf]

No they have very few probability distributions :/