Repository navigation
DOC: Mention take_along_axis in choose #14117
Description
Preamble
In the pymc3 package we implemented a series of distributions. One of such distributions is the Categorical distribution, where we have an ndarray, p, that represents the probabilities of getting one of K categories (where K goes from 0 to p.shape[-1]-1).
Under some circumstances, its good to stack many independent categorical distributions on the same distribution instance (i.e. to write down a multidimensional categorical distribution). In these cases, p.ndim is larger than 1.
PyMC3 focuses on doing MCMC, and to accomplish that, we write down a distribution's log probability. To fix our implementation on multidimensional categorical distributions, we recently switched to use choose instead of advanced indexing (we actually use theano, but theano.tensor.choose later dispatches to numpy.choose). However we encountered the following issue which can be reproduced by the simple following code:
Reproducing code example:
>>> import numpy as np
>>> a = np.random.randint(0, 1000, size=(8, 3, 4))
>>> b = np.random.rand(1000, 3, 4)
>>> np.choose(a, b)
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-8-f1c91ffb39bc> in <module>()
----> 1 np.choose(b, a)
~/anaconda3/lib/python3.6/site-packages/numpy/core/fromnumeric.py in choose(a, choices, out, mode)
420
421 """
--> 422 return _wrapfunc(a, 'choose', choices, out=out, mode=mode)
423
424
~/anaconda3/lib/python3.6/site-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
54 def _wrapfunc(obj, method, *args, **kwds):
55 try:
---> 56 return getattr(obj, method)(*args, **kwds)
57
58 # An AttributeError occurs if the object does not have
ValueError: Need at least 1 and at most 32 array objects.
In [2]: a = np.random.rand(1000, 3, 4)
In [3]: b = np.random.randint(0, 1000, size=(8, 3, 4))
In [4]: np.choose(a, b)
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-11-a1dc619e7b2a> in <module>()
----> 1 np.choose(a, b)
~/anaconda3/lib/python3.6/site-packages/numpy/core/fromnumeric.py in choose(a, choices, out, mode)
420
421 """
--> 422 return _wrapfunc(a, 'choose', choices, out=out, mode=mode)
423
424
~/anaconda3/lib/python3.6/site-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
54 def _wrapfunc(obj, method, *args, **kwds):
55 try:
---> 56 return getattr(obj, method)(*args, **kwds)
57
58 # An AttributeError occurs if the object does not have
ValueError: Need at least 1 and at most 32 array objects.Feature request
It would be really great if this 32 array objects limit were lifted only in the case in which the supplied choices parameter was an ndarray instance and had a non-object dtype (i.e. float64, int64 or others like those).
Numpy/Python version information:
python version: 3.6.7 | packaged by conda-forge | (default, Feb 26 2019, 03:50:56) [GCC 7.3.0]
numpy version: 1.16.1
Activity
seberg commentedon Jul 25, 2019
@lucianopaz
choosereally works on multiple array objects, I am not sure we will change that (although, I admit one could add a different branch/fast path for it).I think your need here should be served very well by
np.take_along_axis, as in:I think it would be good to put
take_along_axisin the "See Also" section and mention it prominently.choicediscourages array inputs, andtake_along_axisshould be mentioned explicitly there probably.EDIT: If you need this in a version which is compatible with much older numpy versions, you can write it as advanced indexing.
lucianopaz commentedon Jul 25, 2019
Thanks a lot @seberg!
take_along_axisseems perfect. When you say that it could be implemented with advanced indexing, do you mean that it could be done without relying on for loops? Just reshapes and indexing? If yes, could you give me a small pointer?seberg commentedon Jul 25, 2019
Yes, if you do not mind diving a bit in the code, in
numpy/lib/shape_base.pyyou can find the code fortake_along_axiswhich does exactly this, but you really only would need it if you support NumPy 1.14 or earlier.[-][Feature request] Choose on arrays[/-][+]DOC: Mention `take_along_axis` in `choose`[/+]lucianopaz commentedon Jul 26, 2019
Thanks again @seberg!
take_along axisis exactly what we need but we can't use numpy's implementation directly because we deal withtheanotensors. The advanced indexing implementation you pointed me to seems perfectly portable to theano though.