Where communities thrive

• Join over 1.5M+ people
• Join over 100K+ communities
• Free without limits
Activity
brandonwillard
@brandonwillard:matrix.org
[m]
yeah, this is all very overlapping functionality
same with Elemwise
but, yes, there is no numpy.vectorize-like helper function in Aesara
and I don't recall jax.vmap doing anything particularly special
it seemed quite literally like a numpy.vectorize clone
the effective difference being that the end result is a JITed function
Ricardo Vieira
@ricardov94:matrix.org
[m]
I think one big difference is that numpy.vectorize requires a base function that works only with scalars, whereas vmap can be built on top of tensor functions?
brandonwillard
@brandonwillard:matrix.org
[m]
I think numpy.vectorize can handle non-scalar functions
that's why you can specify a signature with explicit dimensions for each input
yeah, jax.vmap looks like an clone of numpy.vectorize
with a somewhat different interface
e.g. specified in terms of numbers of axes
or labeled axes?
yeah, it looks like their explicit version of numpy.vectorize just constructs vmap calls
brandonwillard
@brandonwillard:matrix.org
[m]
anyway, we could provide a numpy.vectorize helper function
Ricardo Vieira
@ricardov94:matrix.org
[m]
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)
brandonwillard
@brandonwillard:matrix.org
[m]
yeah, it's pretty useful
this kind of logic is essentially what we're doing in the RandomVariable class interface
e.g. with ndim_supp and ndims_params
and Elemwise is constrained to scalar functions, of course
(same with Composite, since it's just an optimization for Elemwise)
Ricardo Vieira
@ricardov94:matrix.org
[m]
And you think we could do something similar in Aesara, without having to resort to Scan loops internally?
brandonwillard
@brandonwillard:matrix.org
[m]
yeah, we could have a generalized Elemwise
but the implementation/computation details are important
remilouf
@remilouf:matrix.org
[m]
This would be very useful for instance to advance several chains in aehmc
brandonwillard
@brandonwillard:matrix.org
[m]
I think there are directly analogous Numba vectorizations
so that backend might be covered easily
and, obviously, we can map to jax.vmap for the JAX backend
and numpy.vectorize/for-loops in the Python case
I just don't want to deal with an implementation in the current C backend
although that would only involve some straightforward for-loops as well
per usual, I've been avoiding these kinds of things until we replace/remove the C backend
brandonwillard
@brandonwillard:matrix.org
[m]
anyway, as a basic graph abstraction it's a good thing
it's a generalization that has been motivated a few times by our RandomVariable work
I think we're actually talking about gufuncs: https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html#
and since Elemwise is equivalent to ufunc, we need an equivalent Op for gufunc
brandonwillard
@brandonwillard:matrix.org
[m]
yeah, we need an issue for that
brandonwillard
@brandonwillard:matrix.org
[m]
perhaps we could call it Blockwise
and at.vectorize would construct those
actually, this probably would be worth doing sooner than later
especially since there are some optimizations we could write
brandonwillard
@brandonwillard:matrix.org
[m]
we could probably replace the size parameter with this and generalize a lot of the RandomVariable logic
(although the actual NumPy/SciPy size parameter would be used when computed)
we might also want to create a general protocol for specifying signature-like information per Op
because that's valuable information across the board
and, when/if that information isn't available at the Op-level, it can always be provided to a Blockwise
e.g. via the at.vectorize helper/constructor function
Ricardo Vieira
@ricardov94:matrix.org
[m]
Is the C backend somehow in the way of any of this? Or it can happen independently?
brandonwillard
@brandonwillard:matrix.org
[m]
it can happen independently
i.e. we can provide only a Python implementation