Improved handling of boolean masks.

theano/gpuarray/subtensor.py

@@ -4,7 +4,7 @@ import os
 
 import numpy as np
 from six import integer_types
-from six.moves import StringIO
+from six.moves import StringIO, xrange
 
 from theano import tensor, gof, Op
 from theano.gof import ParamsType
@@ -481,6 +481,37 @@ if (err != GA_NO_ERROR) {
         return (0,)
 
 
+def check_and_convert_boolean_masks(input, idx_list):
+    """
+    This function checks if the boolean mask arrays in the index have
+    the right shape and converts them to index arrays by calling nonzero.
+    For each boolean mask, we check if the mask has the
+    same shape as the input. This is enforced in NumPy 0.13.0 and
+    newer, but not by earlier versions. If the size is not the same,
+    this method raises an IndexError.
+    """
+    dim_seen = 0
+    out_idx_list = []
+    for index in idx_list:
+        if index is np.newaxis:
+            # skip, does not count as an input dimension
+            out_idx_list.append(index)
+        elif isinstance(index, np.ndarray) and index.dtype == 'bool':
+            for i in xrange(index.ndim):
+                if index.shape[i] != input.shape[dim_seen + i]:
+                    raise IndexError('boolean index did not match indexed array '
+                                     'along dimension %d; dimension is %d but '
+                                     'corresponding boolean dimension is %d' %
+                                     (dim_seen + i, input.shape[dim_seen + i],
+                                      index.shape[i]))
+            dim_seen += index.ndim
+            out_idx_list += index.nonzero()
+        else:
+            dim_seen += 1
+            out_idx_list.append(index)
+    return out_idx_list
+
+
 class GpuAdvancedSubtensor(HideC, tensor.AdvancedSubtensor):
     """
     AdvancedSubtensor On the GPU.
@@ -499,6 +530,9 @@ class GpuAdvancedSubtensor(HideC, tensor.AdvancedSubtensor):
         x = inputs[0]
         idx = inputs[1:]
 
+        # convert boolean masks to index arrays
+        idx = check_and_convert_boolean_masks(x, idx)
+
         # detect and transpose array indices
         nidx = []
         nshp = list(x.shape)
@@ -631,6 +665,9 @@ class GpuAdvancedIncSubtensor(HideC, tensor.AdvancedIncSubtensor):
             if isinstance(idx[i], gpuarray.GpuArray):
                 idx[i] = np.asarray(idx[i])
 
+        # convert boolean masks to index arrays
+        idx = check_and_convert_boolean_masks(x, idx)
+
         # Insert axes for None indexing
         nidx = []
         nshp = list(x.shape)

theano/tensor/subtensor.py

@@ -2062,7 +2062,7 @@ def as_index_variable(idx):
     if isinstance(idx, gof.Variable) and isinstance(idx.type, NoneTypeT):
         return idx
     idx = theano.tensor.as_tensor_variable(idx)
-    if idx.type.dtype not in theano.tensor.integer_dtypes:
+    if idx.type.dtype not in theano.tensor.integer_dtypes and idx.type.dtype != 'bool':
         raise TypeError('index must be integers')
     return idx
 
@@ -2091,7 +2091,10 @@ def adv_index_broadcastable_pattern(a, idx):
         if isinstance(v.type, SliceType):
             return slice(None, None)
 
-        return np.zeros((2,) * v.ndim, int)
+        if v.dtype == 'bool':
+            return np.ones((2,) * v.ndim, v.dtype)
+        else:
+            return np.zeros((2,) * v.ndim, int)
 
     newidx = tuple(map(replace_slice, idx))
 
@@ -2101,6 +2104,32 @@ def adv_index_broadcastable_pattern(a, idx):
     return tuple([dim == 1 for dim in retshape])
 
 
+def check_advanced_indexing_dimensions(input, idx_list):
+    """
+    This function checks if the index list in idx_list is correct.
+    If there are any boolean masks, we check if the mask has the
+    same shape as the input. This is enforced in NumPy 0.13.0 and
+    newer, but not by earlier versions. If the size is not the same,
+    this method raises an IndexError.
+    """
+    dim_seen = 0
+    for index in idx_list:
+        if index is np.newaxis:
+            # skip, does not count as an input dimension
+            pass
+        elif isinstance(index, np.ndarray) and index.dtype == 'bool':
+            for i in xrange(index.ndim):
+                if index.shape[i] != input.shape[dim_seen + i]:
+                    raise IndexError('boolean index did not match indexed array '
+                                     'along dimension %d; dimension is %d but '
+                                     'corresponding boolean dimension is %d' %
+                                     (dim_seen + i, input.shape[dim_seen + i],
+                                      index.shape[i]))
+            dim_seen += index.ndim
+        else:
+            dim_seen += 1
+
+
 class AdvancedSubtensor(Op):
     """
     Return a subtensor copy, using advanced indexing.
@@ -2146,6 +2175,7 @@ class AdvancedSubtensor(Op):
 
     def perform(self, node, inputs, out_):
         out, = out_
+        check_advanced_indexing_dimensions(inputs[0], inputs[1:])
         rval = inputs[0].__getitem__(inputs[1:])
         # When there are no arrays, we are not actually doing advanced
         # indexing, so __getitem__ will not return a copy.
@@ -2215,6 +2245,8 @@ class AdvancedIncSubtensor(Op):
         # TODO: 1. opt to make this in place 2. generalize as described in
         # AdvancedSubtensor's perform TODO
 
+        check_advanced_indexing_dimensions(inputs[0], inputs[2:])
+
         out, = out_
         if not self.inplace:
             out[0] = inputs[0].copy()

theano/tensor/tests/test_subtensor.py

@@ -363,6 +363,11 @@ class T_subtensor(unittest.TestCase, utt.TestOptimizationMixin):
             assert_array_equal(tval, numpy_tval)
 
     def test_boolean(self):
+        def numpy_inc_subtensor(x, idx, a):
+            x = x.copy()
+            x[idx] += a
+            return x
+
         numpy_n = np.arange(6, dtype=self.dtype).reshape((2, 3))
         n = self.shared(numpy_n)
 
@@ -374,6 +379,10 @@ class T_subtensor(unittest.TestCase, utt.TestOptimizationMixin):
         # indexing with a mask for some dimensions
         mask = np.array([True, False])
         assert_array_equal(numpy_n[mask], n[mask].eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n, mask, 1),
+                           inc_subtensor(n[mask], 1).eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n, mask, numpy_n[mask]),
+                           inc_subtensor(n[mask], n[mask]).eval())
 
         # indexing with a mask for the second dimension
         mask = np.array([True, False, True])
@@ -382,41 +391,66 @@ class T_subtensor(unittest.TestCase, utt.TestOptimizationMixin):
         assert_array_equal(numpy_n[:, mask], n[:, self.shared(mask)].eval())
         assert_array_equal(numpy_n[1:, mask], n[1:, mask].eval())
         assert_array_equal(numpy_n[:1, mask], n[:1, mask].eval())
+        assert_array_equal(numpy_n[1:, mask, np.newaxis], n[1:, mask, np.newaxis].eval())
+        assert_array_equal(numpy_n[np.newaxis, 1:, mask], n[np.newaxis, 1:, mask].eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n, [0, mask], 1),
+                           inc_subtensor(n[0, mask], 1).eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n, [Ellipsis, mask], 1),
+                           inc_subtensor(n[:, mask], 1).eval())
 
         # indexing with a boolean ndarray
         mask = np.array([[True, False, True], [False, False, True]])
         assert_array_equal(numpy_n[mask], n[mask].eval())
         assert_array_equal(numpy_n[mask], n[self.shared(mask)].eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n, mask, 1),
+                           inc_subtensor(n[mask], 1).eval())
 
         # indexing with ellipsis
         numpy_n4 = np.arange(48, dtype=self.dtype).reshape((2, 3, 4, 2))
         n4 = self.shared(numpy_n4)
         assert_array_equal(numpy_n4[numpy_n > 2, ...], n4[n > 2, ...].eval())
         assert_array_equal(numpy_n4[numpy_n > 2, ..., 1], n4[n > 2, ..., 1].eval())
+        assert_array_equal(numpy_n4[numpy_n > 2, ..., 0, 1], n4[n > 2, ..., 0, 1].eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n4, [numpy_n > 2, Ellipsis], 1),
+                           inc_subtensor(n4[n > 2, ...], 1).eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n4, [numpy_n > 2, Ellipsis, 1], 1),
+                           inc_subtensor(n4[n > 2, ..., 1], 1).eval())
+        assert_array_equal(numpy_inc_subtensor(numpy_n4, [numpy_n > 2, Ellipsis, 0, 1], 1),
+                           inc_subtensor(n4[n > 2, ..., 0, 1], 1).eval())
 
         # the boolean mask should have the correct shape
         # - too large, padded with True
         mask = np.array([True, False, True])
         self.assertRaises(IndexError, n[mask].eval)
         self.assertRaises(IndexError, n[mask, ...].eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask], 1).eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask, ...], 1).eval)
         mask = np.array([[True, False, False, True], [False, True, False, True]])
         self.assertRaises(IndexError, n[mask].eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask], 1).eval)
         # - too large, padded with False (this works in NumPy < 0.13.0)
         mask = np.array([True, False, False])
         self.assertRaises(IndexError, n[mask].eval)
         self.assertRaises(IndexError, n[mask, ...].eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask], 1).eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask, ...], 1).eval)
         mask = np.array([[True, False, False, False], [False, True, False, False]])
         self.assertRaises(IndexError, n[mask].eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask], 1).eval)
         # - mask too small (this works in NumPy < 0.13.0)
         mask = np.array([True])
         self.assertRaises(IndexError, n[mask].eval)
         self.assertRaises(IndexError, n[mask, ...].eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask], 1).eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask, ...], 1).eval)
         mask = np.array([[True], [True]])
         self.assertRaises(IndexError, n[mask].eval)
+        self.assertRaises(IndexError, inc_subtensor(n[mask], 1).eval)
         # - too many dimensions
         mask = np.array([[[True, False, False],
                           [False, True, False]]])
-        self.assertRaises(IndexError, n[mask].eval)
+        self.assertRaises(IndexError, n.__getitem__, mask)
+        self.assertRaises(IndexError, n.__getitem__, mask)
 
         # special cases: Python bools and bools nested in Python arrays are not supported
         self.assertRaises(TypeError, n.__getitem__, (True,))

theano/tensor/var.py

@@ -477,15 +477,19 @@ class _tensor_py_operators(object):
         elif not isinstance(args, tuple):
             args = args,
 
-        # Convert boolean arrays to calls to mask.nonzero()
-        tmp_args = []
-        for arg in args:
-            # NumPy arrays or tensors of type bool can be converted to
-            # normal integer indices.
-            if (isinstance(arg, (np.ndarray, theano.tensor.Variable)) and
-                    hasattr(arg, 'dtype') and hasattr(arg, 'nonzero') and
-                    arg.dtype == 'bool'):
-                tmp_args += arg.nonzero()
+        # Count the dimensions, check for bools and find ellipses.
+        ellipses = []
+        index_dim_count = 0
+        for i, arg in enumerate(args):
+            if arg is np.newaxis:
+                # no increase in index_dim_count
+                pass
+            elif arg is Ellipsis:
+                # no increase in index_dim_count
+                ellipses.append(i)
+            elif (isinstance(arg, (np.ndarray, theano.tensor.Variable)) and
+                    hasattr(arg, 'dtype') and arg.dtype == 'bool'):
+                index_dim_count += arg.ndim
             else:
                 # Python arrays can contain a mixture of bools and integers,
                 # which requires complex rules to handle all special cases.
@@ -499,25 +503,22 @@ class _tensor_py_operators(object):
                                     'To use a boolean mask, convert the mask to '
                                     'a NumPy array first, e.g., '
                                     'tensor[numpy.array([True, False])].')
-                tmp_args.append(arg)
-        args = tuple(tmp_args)
+                index_dim_count += 1
+
+        # Check if the number of dimensions isn't too large.
+        if self.ndim < index_dim_count:
+            raise IndexError('too many indices for array')
 
         # Convert an Ellipsis if provided into an appropriate number of
         # slice(None).
-        ellipses = [i
-                    for i, index in enumerate(args)
-                    if index is Ellipsis]
         if len(ellipses) > 1:
             raise IndexError(
                 "an index can only have a single Ellipsis (`...`)")
         elif len(ellipses) == 1:
-            new_axes = sum(1
-                           for index in args
-                           if index is np.newaxis)  # numpy.newaxis is None
             ellipsis_at = ellipses[0]
             args = list(args)
             args[ellipsis_at: ellipsis_at + 1] = (
-                [slice(None)] * (self.ndim - (len(args) - 1 - new_axes)))
+                [slice(None)] * (self.ndim - index_dim_count))
 
         # Force input to be int64 datatype if input is an empty list or tuple
         # Else leave it as is if it is a real number
@@ -533,7 +534,12 @@ class _tensor_py_operators(object):
         axis = None
         for i, arg in enumerate(args):
             try:
-                if arg is not np.newaxis:
+                if (isinstance(arg, (np.ndarray, theano.tensor.Variable)) and
+                        hasattr(arg, 'dtype') and arg.dtype == 'bool'):
+                    advanced = True
+                    axis = None
+                    break
+                elif arg is not np.newaxis:
                     theano.tensor.subtensor.Subtensor.convert(arg)
             except theano.tensor.subtensor.AdvancedIndexingError:
                 if advanced: