Merge pull request #2 from lamblin/take-op-c-code-clean

theano/tensor/basic.py

@@ -6547,10 +6547,15 @@ class AdvancedSubtensor1(Op):
         # If i.dtype is more precise than numpy.intp (int32 on 32-bit machines,
         # int64 on 64-bit machines), numpy may raise the following error:
         # TypeError: array cannot be safely cast to required type.
-        # Since we will probably not have an array with more than 2**31 items
-        # on a 32-bit arch, I suppose it is safe to cast i into intp.
+        # We need to check if values in i can fit in numpy.intp, because
+        # if they don't, that should be an error (no array can have that
+        # many elements on a 32-bit arch).
         if i.dtype != numpy.intp:
-            i = theano._asarray(i, dtype=numpy.intp)
+            i_ = theano._asarray(i, dtype=numpy.intp)
+            if numpy.any(i != i_):
+                raise IndexError('index contains values that are bigger than '
+                        'the maximum array size on this system.', i)
+            i = i_
 
         out[0] = x.take(i, axis=0, out=o)
 
@@ -6596,8 +6601,24 @@ class AdvancedSubtensor1(Op):
         return """
             PyObject *indices;
             if (PyArray_TYPE(%(i_name)s) != NPY_INTP) {
-                // This cast makes c_code mimic the logic in perform(), but
-                // also makes theano code less safe.
+                // Cast %(i_name)s to NPY_INTP (expected by PyArray_TakeFrom),
+                // if all values fit.
+                PyObject* py_min_val, py_max_val;
+                npy_int64 min_val, max_val;
+                py_min_val = PyArray_Min(%(i_name)s, NPY_MAXDIMS, min_val);
+                py_max_val = PyArray_Max(%(i_name)s, NPY_MAXDIMS, max_val);
+                min_val = PyLong_AsLongLong(py_min_val);
+                max_val = PyLong_AsLongLong(py_max_val);
+                Py_CLEAR(py_min_val);
+                Py_CLEAR(py_max_val);
+
+                if ((min_val < NPY_MIN_INTP) || (max_val > NPY_MAX_INTP))
+                {
+                    PyExc_SetErr(PyExc_IndexError, "Index contains values "
+                                 "that are bigger than the maximum array "
+                                 "size on this system.");
+                    %(fail)s;
+                }
                 indices = PyArray_Cast(%(i_name)s, NPY_INTP);
                 if (indices == NULL) {
                     %(fail)s;
@@ -6638,7 +6659,7 @@ class AdvancedSubtensor1(Op):
         """ % locals()
 
     def c_code_cache_version(self):
-        return (0, 0, 3)
+        return (0, 0, 4)
 
 advanced_subtensor1 = AdvancedSubtensor1()