Elemwise
numpy.vectorize
-like helper function in Aesara
jax.vmap
doing anything particularly special
numpy.vectorize
clone
numpy.vectorize
can handle non-scalar functions
signature
with explicit dimensions for each input
jax.vmap
looks like an clone of numpy.vectorize
numpy.vectorize
just constructs vmap
calls
numpy.vectorize
helper function
signature
keyword:import numpy as np
def dirichlet_logp(a, x):
return st.dirichlet(a).logpdf(x)
vfunc = np.vectorize(dirichlet_logp, signature='(n),(n)->()')
vfunc(np.arange(1, 10).reshape(3, 3), np.ones(3)/3)
RandomVariable
class interface
ndim_supp
and ndims_params
Elemwise
is constrained to scalar functions, of course
Composite
, since it's just an optimization for Elemwise
)
aehmc
jax.vmap
for the JAX backend
numpy.vectorize
/for
-loops in the Python case
for
-loops as well
RandomVariable
work
gufunc
s: https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html#
Elemwise
is equivalent to ufunc
, we need an equivalent Op
for gufunc
at.vectorize
would construct those
size
parameter with this and generalize a lot of the RandomVariable
logic
size
parameter would be used when computed)
signature
-like information per Op
Op
-level, it can always be provided to a Blockwise
at.vectorize
helper/constructor function