Add handling of negative axis for Join and GpuJoin

theano/sandbox/cuda/basic_ops.py

@@ -3041,24 +3041,6 @@ class GpuJoin(tensor.Join, GpuOp):
         as_tensor_variable_args = [as_cuda_ndarray_variable(x)
                                    for x in tensors]
 
-        # Get joining axis as int
-        axis_int = 0
-        if not isinstance(axis, int):
-            try:
-                # Note : `get_scalar_constant_value` returns a ndarray not
-                # an int
-                axis_int = int(tensor.get_scalar_constant_value(axis))
-
-            except tensor.basic.NotScalarConstantError:
-                pass
-        else:
-            axis_int = axis
-
-        if (axis_int < 0):
-            # Since all tensors must have the same number of dimensions,
-            # we simply add the number of dimensions for the first tensor
-            axis = axis + as_tensor_variable_args[0].ndim
-
         output_maker = \
                 lambda bcast: CudaNdarrayType(broadcastable=bcast)()
 
@@ -3071,6 +3053,12 @@ class GpuJoin(tensor.Join, GpuOp):
         axis, cndas = axis_and_tensors[0], axis_and_tensors[1:]
         # In case axis is numpy.int8 and has no __index__() method
         axis = int(axis)
+        ndim = tensors[0].ndim
+        if axis < -ndim:
+            raise IndexError("Join axis %d out of bounds [0, %d)" %
+                             (axis, ndim))
+        if axis < 0:
+            axis += ndim
 
         # compute size/shape
         width_sum = 0
@@ -3134,7 +3122,7 @@ class GpuJoin(tensor.Join, GpuOp):
 
         # getting the shapes of all the involved tensors (input[0]+out)
         str = """
-        const int axis = PyInt_AsLong((PyObject*)%(axis)s);
+        int axis = PyInt_AsLong((PyObject*)%(axis)s);
         const int nd = %(nd)s;
         int shape_out[nd];
         int width_sum = 0;
@@ -3150,9 +3138,22 @@ class GpuJoin(tensor.Join, GpuOp):
 
         """ % locals()
 
+        # Test negative axis
+        str += """
+        if( axis < -nd ){
+            PyErr_Format(PyExc_IndexError, 
+                         "Join axis %%d out of bounds [0, %%d)", axis, nd);
+            %(fail)s
+        }
+
+        if( axis < 0 ){
+            axis = axis + nd;
+        }
+        """ % locals()
+
         # getting the shapes of all the involved tensors (input[1:])
         # + check: all input tensors have same shape as final out
-        # execept for "axis" dimension
+        # except for "axis" dimension
         # shape_%(cdna)s[nd] is initialized before, to prevent following
         # error: jump to label __label_9 crosses initialization of
         # shape_%(cdna)s[nd]
@@ -3266,7 +3267,7 @@ class GpuJoin(tensor.Join, GpuOp):
         return str
 
     def c_code_cache_version(self):
-        return (5,)
+        return (6,)
 
 gpu_join = GpuJoin()
 

theano/tensor/basic.py

@@ -3453,10 +3453,18 @@ class Join(Op):
                 # that broadcastable flag was False had length 1 along
                 # this dimension, and therefore this dimension should
                 # be broadcastable for the output.
+
+                if axis < -ndim:
+                    raise IndexError("Join axis %d out of bounds [0, %d)" %
+                                     (axis, ndim))
+                if axis < 0:
+                    axis += ndim
+
                 for x in as_tensor_variable_args:
                     for current_axis, bflag in enumerate(x.type.broadcastable):
-                        # This Op supports negative axes, so only consider modulo
-                        if current_axis == axis % ndim:
+                        # Constant negative axis can no longer be negative at 
+                        # this point. It safe to compare this way.
+                        if current_axis == axis:
                             continue
                         if bflag:
                             bcastable[current_axis] = True
@@ -3489,14 +3497,20 @@ class Join(Op):
     def perform(self, node, axis_and_tensors, out_):
         out, = out_
         axis, tensors = axis_and_tensors[0], axis_and_tensors[1:]
+        ndim = tensors[0].ndim
+        if axis < -ndim:
+            raise IndexError("Join axis %d out of bounds [0, %d)" %
+                             (axis, ndim))
+
         out[0] = theano._asarray(numpy.concatenate(tensors, axis=axis),
                                  dtype=node.outputs[0].type.dtype)
 
     def c_code_cache_version(self):
-        return (2,)
+        return (3,)
 
     def c_code(self, node, name, inputs, outputs, sub):
         axis, tensors = inputs[0], inputs[1:]
+        input_1 = tensors[0]
         l = len(tensors)
         out, = outputs
         fail = sub['fail']
@@ -3511,9 +3525,16 @@ class Join(Op):
             """ % locals()
         code += """
         //PyObject* PyArray_Concatenate(PyObject* obj, int axis)
+        int axis = ((%(adtype)s *)PyArray_DATA(%(axis)s))[0];
+        int ndim = PyArray_NDIM(%(input_1)s);
+        if( axis < -ndim ){
+            PyErr_Format(PyExc_IndexError, 
+                         "Join axis %%d out of bounds [0, %%d)", axis, ndim);
+            %(fail)s
+        }
+
         Py_XDECREF(%(out)s);
-        %(out)s = (PyArrayObject *)PyArray_Concatenate(list,
-                      ((%(adtype)s *)PyArray_DATA(%(axis)s))[0]);
+        %(out)s = (PyArrayObject *)PyArray_Concatenate(list, axis);
 
         Py_DECREF(list);
         if(!%(out)s){

theano/tensor/tests/test_basic.py

@@ -3564,8 +3564,8 @@ class T_Join_and_Split(unittest.TestCase):
     def test_join_matrixV(self):
         """variable join axis"""
         v = numpy.array([[.1, .2, .3], [.4, .5, .6]], dtype=self.floatX)
-        a = self.shared(v.copy())
-        b = as_tensor_variable(v.copy())
+        a = self.shared(v)
+        b = as_tensor_variable(v)
         ax = lscalar()
         s = join(ax, a, b)
 
@@ -3587,6 +3587,75 @@ class T_Join_and_Split(unittest.TestCase):
         utt.verify_grad(lambda a, b: join(0, a, b), [v, 2 * v], mode=self.mode)
         utt.verify_grad(lambda a, b: join(1, a, b), [v, 2 * v], mode=self.mode)
 
+    def test_join_matrixV_negative_axis(self):
+        """variable join negative axis"""
+        v = numpy.array([[.1, .2, .3], [.4, .5, .6]], dtype=self.floatX)
+        a = self.shared(v)
+        b = as_tensor_variable(v)
+        ax = lscalar()
+        s = join(ax, a, b)
+
+        f = inplace_func([ax], [s], mode=self.mode)
+        topo = f.maker.fgraph.toposort()
+        assert [True for node in topo
+                if isinstance(node.op, type(self.join_op))]
+
+        want = numpy.array([[.1, .2, .3, .1, .2, .3],
+                            [.4, .5, .6, .4, .5, .6]])
+
+        got = f(-1)
+        assert numpy.allclose(got, want)
+
+        want = numpy.array([[.1, .2, .3], [.4, .5, .6],
+                            [.1, .2, .3], [.4, .5, .6]])
+        got = f(-2)
+        assert numpy.allclose(got, want)
+
+        try:
+            got = f(-3)
+            assert False
+        except IndexError:
+            pass
+
+    def test_join_matrixC_negative_axis(self):
+        """constant join negative axis"""
+        v = numpy.array([[.1, .2, .3], [.4, .5, .6]], dtype=self.floatX)
+        a = self.shared(v)
+        b = as_tensor_variable(v)
+
+        s = join(-1, a, b)
+        f = theano.function([], [s], mode=self.mode)
+        topo = f.maker.fgraph.toposort()
+        assert [True for node in topo
+                if isinstance(node.op, type(self.join_op))]
+
+        want = numpy.array([[.1, .2, .3, .1, .2, .3],
+                            [.4, .5, .6, .4, .5, .6]])
+
+        got = f()
+        assert numpy.allclose(got, want)
+
+        s = join(-2, a, b)
+        f = theano.function([], [s], mode=self.mode)
+        topo = f.maker.fgraph.toposort()
+        assert [True for node in topo
+                if isinstance(node.op, type(self.join_op))]
+
+        want = numpy.array([[.1, .2, .3], [.4, .5, .6],
+                            [.1, .2, .3], [.4, .5, .6]])
+
+        got = f()
+        assert numpy.allclose(got, want)
+
+        try:
+            s = join(-3, a, b)
+            assert False
+        except IndexError:
+            pass
+
+        utt.verify_grad(lambda a, b: join(-1, a, b), [v, 2 * v],
+                        mode=self.mode)
+
     def test_vector_len(self):
         x = lscalar('x')
         y = dscalar('y')