Remove take_along_axis in favor of Aesara's implementation

Author: ricardoV94

pymc/aesaraf.py

@@ -57,7 +57,7 @@
 from aesara.tensor.var import TensorConstant, TensorVariable
 
 from pymc.exceptions import ShapeError
-from pymc.vartypes import continuous_types, int_types, isgenerator, typefilter
+from pymc.vartypes import continuous_types, isgenerator, typefilter
 
 PotentialShapeType = Union[
     int, np.ndarray, Tuple[Union[int, Variable], ...], List[Union[int, Variable]], Variable
@@ -80,7 +80,6 @@
     "generator",
     "set_at_rng",
     "at_rng",
-    "take_along_axis",
     "convert_observed_data",
 ]
 
@@ -854,61 +853,6 @@ def largest_common_dtype(tensors):
     return np.stack([np.ones((), dtype=dtype) for dtype in dtypes]).dtype
 
 
-def _make_along_axis_idx(arr_shape, indices, axis):
-    # compute dimensions to iterate over
-    if str(indices.dtype) not in int_types:
-        raise IndexError("`indices` must be an integer array")
-    shape_ones = (1,) * indices.ndim
-    dest_dims = list(range(axis)) + [None] + list(range(axis + 1, indices.ndim))
-
-    # build a fancy index, consisting of orthogonal aranges, with the
-    # requested index inserted at the right location
-    fancy_index = []
-    for dim, n in zip(dest_dims, arr_shape):
-        if dim is None:
-            fancy_index.append(indices)
-        else:
-            ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1 :]
-            fancy_index.append(at.arange(n).reshape(ind_shape))
-
-    return tuple(fancy_index)
-
-
-def take_along_axis(arr, indices, axis=0):
-    """Take values from the input array by matching 1d index and data slices.
-
-    This iterates over matching 1d slices oriented along the specified axis in
-    the index and data arrays, and uses the former to look up values in the
-    latter. These slices can be different lengths.
-
-    Functions returning an index along an axis, like argsort and argpartition,
-    produce suitable indices for this function.
-    """
-    arr = at.as_tensor_variable(arr)
-    indices = at.as_tensor_variable(indices)
-    # normalize inputs
-    if axis is None:
-        arr = arr.flatten()
-        arr_shape = (len(arr),)  # flatiter has no .shape
-        _axis = 0
-    else:
-        if axis < 0:
-            _axis = arr.ndim + axis
-        else:
-            _axis = axis
-        if _axis < 0 or _axis >= arr.ndim:
-            raise ValueError(
-                "Supplied `axis` value {} is out of bounds of an array with "
-                "ndim = {}".format(axis, arr.ndim)
-            )
-        arr_shape = arr.shape
-    if arr.ndim != indices.ndim:
-        raise ValueError("`indices` and `arr` must have the same number of dimensions")
-
-    # use the fancy index
-    return arr[_make_along_axis_idx(arr_shape, indices, _axis)]
-
-
 @local_optimizer(tracks=[CheckParameterValue])
 def local_remove_check_parameter(fgraph, node):
     """Rewrite that removes Aeppl's CheckParameterValue

pymc/distributions/discrete.py

@@ -31,7 +31,7 @@
 
 import pymc as pm
 
-from pymc.aesaraf import floatX, intX, take_along_axis
+from pymc.aesaraf import floatX, intX
 from pymc.distributions.dist_math import (
     betaln,
     binomln,
@@ -1318,7 +1318,7 @@ def logp(value, p):
                 p = at.shape_padleft(p, value_clip.ndim - p_.ndim)
             pattern = (p.ndim - 1,) + tuple(range(p.ndim - 1))
             a = at.log(
-                take_along_axis(
+                at.take_along_axis(
                     p.dimshuffle(pattern),
                     value_clip,
                 )

pymc/tests/test_aesaraf.py

@@ -12,8 +12,6 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-from itertools import product
-
 import aesara
 import aesara.tensor as at
 import numpy as np
@@ -29,19 +27,16 @@
 from aesara.tensor.random.basic import normal, uniform
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
-from aesara.tensor.type import TensorType
 from aesara.tensor.var import TensorVariable
 
 import pymc as pm
 
 from pymc.aesaraf import (
-    _conversion_map,
     change_rv_size,
     compile_pymc,
     convert_observed_data,
     extract_obs_data,
     rvs_to_value_vars,
-    take_along_axis,
     walk_model,
 )
 from pymc.distributions.dist_math import check_parameters
@@ -166,198 +161,6 @@ def _make_along_axis_idx(arr_shape, indices, axis):
     return tuple(fancy_index)
 
 
-if hasattr(np, "take_along_axis"):
-    np_take_along_axis = np.take_along_axis
-else:
-
-    def np_take_along_axis(arr, indices, axis):
-        if arr.shape[axis] <= 32:
-            # We can safely test with numpy's choose
-            arr = np.moveaxis(arr, axis, 0)
-            indices = np.moveaxis(indices, axis, 0)
-            out = np.choose(indices, arr)
-            return np.moveaxis(out, 0, axis)
-        else:
-            # numpy's choose cannot handle such a large axis so we
-            # just use the implementation of take_along_axis. This is kind of
-            # cheating because our implementation is the same as the one below
-            if axis < 0:
-                _axis = arr.ndim + axis
-            else:
-                _axis = axis
-            if _axis < 0 or _axis >= arr.ndim:
-                raise ValueError(f"Supplied axis {axis} is out of bounds")
-            return arr[_make_along_axis_idx(arr.shape, indices, _axis)]
-
-
-class TestTakeAlongAxis:
-    def setup_class(self):
-        self.inputs_buffer = dict()
-        self.output_buffer = dict()
-        self.func_buffer = dict()
-
-    def _input_tensors(self, shape, floatX):
-        intX = str(_conversion_map[floatX])
-        ndim = len(shape)
-        arr = TensorType(floatX, [False] * ndim)("arr")
-        indices = TensorType(intX, [False] * ndim)("indices")
-        arr.tag.test_value = np.zeros(shape, dtype=floatX)
-        indices.tag.test_value = np.zeros(shape, dtype=intX)
-        return arr, indices
-
-    def get_input_tensors(self, shape, floatX):
-        ndim = len(shape)
-        try:
-            return self.inputs_buffer[(ndim, floatX)]
-        except KeyError:
-            arr, indices = self._input_tensors(shape, floatX)
-            self.inputs_buffer[(ndim, floatX)] = arr, indices
-            return arr, indices
-
-    def _output_tensor(self, arr, indices, axis):
-        return take_along_axis(arr, indices, axis)
-
-    def get_output_tensors(self, shape, axis, floatX):
-        ndim = len(shape)
-        try:
-            return self.output_buffer[(ndim, axis, floatX)]
-        except KeyError:
-            arr, indices = self.get_input_tensors(shape, floatX)
-            out = self._output_tensor(arr, indices, axis)
-            self.output_buffer[(ndim, axis, floatX)] = out
-            return out
-
-    def _function(self, arr, indices, out):
-        return aesara.function([arr, indices], [out])
-
-    def get_function(self, shape, axis, floatX):
-        ndim = len(shape)
-        try:
-            return self.func_buffer[(ndim, axis, floatX)]
-        except KeyError:
-            arr, indices = self.get_input_tensors(shape, floatX)
-            out = self.get_output_tensors(shape, axis, floatX)
-            func = self._function(arr, indices, out)
-            self.func_buffer[(ndim, axis, floatX)] = func
-            return func
-
-    @staticmethod
-    def get_input_values(shape, axis, samples, floatX):
-        intX = str(_conversion_map[floatX])
-        arr = np.random.randn(*shape).astype(floatX)
-        size = list(shape)
-        size[axis] = samples
-        size = tuple(size)
-        indices = np.random.randint(low=0, high=shape[axis], size=size, dtype=intX)
-        return arr, indices
-
-    @pytest.mark.parametrize(
-        ["shape", "axis", "samples"],
-        product(
-            [
-                (1,),
-                (3,),
-                (3, 1),
-                (3, 2),
-                (1, 1),
-                (1, 2),
-                (40, 40),  # choose fails here
-                (5, 1, 1),
-                (5, 1, 2),
-                (5, 3, 1),
-                (5, 3, 2),
-            ],
-            [0, -1],
-            [1, 10],
-        ),
-        ids=str,
-    )
-    @pytest.mark.parametrize("floatX", ["float32", "float64"])
-    def test_take_along_axis(self, shape, axis, samples, floatX):
-        with aesara.config.change_flags(floatX=floatX):
-            arr, indices = self.get_input_values(shape, axis, samples, floatX)
-            func = self.get_function(shape, axis, floatX)
-            assert np.allclose(np_take_along_axis(arr, indices, axis=axis), func(arr, indices)[0])
-
-    @pytest.mark.parametrize(
-        ["shape", "axis", "samples"],
-        product(
-            [
-                (1,),
-                (3,),
-                (3, 1),
-                (3, 2),
-                (1, 1),
-                (1, 2),
-                (40, 40),  # choose fails here
-                (5, 1, 1),
-                (5, 1, 2),
-                (5, 3, 1),
-                (5, 3, 2),
-            ],
-            [0, -1],
-            [1, 10],
-        ),
-        ids=str,
-    )
-    @pytest.mark.parametrize("floatX", ["float32", "float64"])
-    def test_take_along_axis_grad(self, shape, axis, samples, floatX):
-        with aesara.config.change_flags(floatX=floatX):
-            if axis < 0:
-                _axis = len(shape) + axis
-            else:
-                _axis = axis
-            # Setup the aesara function
-            t_arr, t_indices = self.get_input_tensors(shape, floatX)
-            t_out2 = aesara.grad(
-                at.sum(self._output_tensor(t_arr**2, t_indices, axis)),
-                t_arr,
-            )
-            func = aesara.function([t_arr, t_indices], [t_out2])
-
-            # Test that the gradient gives the same output as what is expected
-            arr, indices = self.get_input_values(shape, axis, samples, floatX)
-            expected_grad = np.zeros_like(arr)
-            slicer = [slice(None)] * len(shape)
-            for i in range(indices.shape[axis]):
-                slicer[axis] = i
-                inds = indices[tuple(slicer)].reshape(shape[:_axis] + (1,) + shape[_axis + 1 :])
-                inds = _make_along_axis_idx(shape, inds, _axis)
-                expected_grad[inds] += 1
-            expected_grad *= 2 * arr
-            out = func(arr, indices)[0]
-            assert np.allclose(out, expected_grad)
-
-    @pytest.mark.parametrize("axis", [-4, 4], ids=str)
-    @pytest.mark.parametrize("floatX", ["float32", "float64"])
-    def test_axis_failure(self, axis, floatX):
-        with aesara.config.change_flags(floatX=floatX):
-            arr, indices = self.get_input_tensors((3, 1), floatX)
-            with pytest.raises(ValueError):
-                take_along_axis(arr, indices, axis=axis)
-
-    @pytest.mark.parametrize("floatX", ["float32", "float64"])
-    def test_ndim_failure(self, floatX):
-        with aesara.config.change_flags(floatX=floatX):
-            intX = str(_conversion_map[floatX])
-            arr = TensorType(floatX, [False] * 3)("arr")
-            indices = TensorType(intX, [False] * 2)("indices")
-            arr.tag.test_value = np.zeros((1,) * arr.ndim, dtype=floatX)
-            indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=intX)
-            with pytest.raises(ValueError):
-                take_along_axis(arr, indices)
-
-    @pytest.mark.parametrize("floatX", ["float32", "float64"])
-    def test_dtype_failure(self, floatX):
-        with aesara.config.change_flags(floatX=floatX):
-            arr = TensorType(floatX, [False] * 3)("arr")
-            indices = TensorType(floatX, [False] * 3)("indices")
-            arr.tag.test_value = np.zeros((1,) * arr.ndim, dtype=floatX)
-            indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=floatX)
-            with pytest.raises(IndexError):
-                take_along_axis(arr, indices)
-
-
 def test_extract_obs_data():
 
     with pytest.raises(TypeError):