Merge pull request #5349 from gvtulder/f-softmax-zero-length

theano/gpuarray/dnn_softmax.c

@@ -39,14 +39,19 @@ int APPLY_SPECIFIC(softmax)(PyGpuArrayObject *x,
   PyGpuContextObject *c = x->context;
   cudnnStatus_t err;
 
-  if (c_set_tensorNd(x, APPLY_SPECIFIC(input)) != 0)
-    return 1;
-
   if (theano_prep_output(out, PyGpuArray_NDIM(x),
                          PyGpuArray_DIMS(x), x->ga.typecode,
                          GA_C_ORDER, c) != 0)
     return 1;
 
+  // Directly return the output if any of the dimensions is 0.
+  // (cuDNN does not support zero-length dimensions.)
+  if (PyGpuArray_SIZE(*out) == 0)
+    return 0;
+
+  if (c_set_tensorNd(x, APPLY_SPECIFIC(input)) != 0)
+    return 1;
+
   if (c_set_tensorNd(*out, APPLY_SPECIFIC(output)) != 0)
     return 1;
 

theano/gpuarray/dnn_softmax_grad.c

@@ -50,16 +50,21 @@ int APPLY_SPECIFIC(softmax_grad)(PyGpuArrayObject *dy,
   PyGpuContextObject *c = dy->context;
   cudnnStatus_t err;
 
-  if (c_set_tensorNd(dy, APPLY_SPECIFIC(dy)) != 0)
-    return 1;
-  if (c_set_tensorNd(sm, APPLY_SPECIFIC(sm)) != 0)
-    return 1;
-
   if (theano_prep_output(dx, PyGpuArray_NDIM(dy),
                          PyGpuArray_DIMS(dy), dy->ga.typecode,
                          GA_C_ORDER, c) != 0)
     return 1;
 
+  // Directly return the output if any of the dimensions is 0.
+  // (cuDNN does not support zero-length dimensions.)
+  if (PyGpuArray_SIZE(*dx) == 0)
+    return 0;
+
+  if (c_set_tensorNd(dy, APPLY_SPECIFIC(dy)) != 0)
+    return 1;
+  if (c_set_tensorNd(sm, APPLY_SPECIFIC(sm)) != 0)
+    return 1;
+
   if (c_set_tensorNd(*dx, APPLY_SPECIFIC(dx)) != 0)
     return 1;
 

theano/gpuarray/tests/test_dnn.py

@@ -1131,7 +1131,24 @@ class test_SoftMax(test_nnet.test_SoftMax):
             raise SkipTest(dnn.dnn_available.msg)
 
     def test_softmax_shape_0(self):
-        raise SkipTest("Cudnn doesn't support 0 shapes")
+        dims = (2, 0, 4, 5)
+        data = numpy.arange(
+            numpy.product(dims),
+            dtype=theano.config.floatX
+        ).reshape(dims)
+
+        # Verify the forward op
+        x_gpu = T.tensor4('x_gpu')
+        f_gpu = dnn.GpuDnnSoftmax('accurate', 'channel')(x_gpu)
+        f_gpu = theano.function([x_gpu], f_gpu, mode=self.mode)
+        assert f_gpu(data).shape == dims
+
+        # Verify the gradient op
+        dy_gpu = T.tensor4('dy_gpu')
+        sm_gpu = T.tensor4('sm_gpu')
+        f_grad_gpu = dnn.GpuDnnSoftmaxGrad('accurate', 'channel')(dy_gpu, sm_gpu)
+        f_grad_gpu = theano.function([dy_gpu, sm_gpu], f_grad_gpu, mode=self.mode)
+        assert f_grad_gpu(data, data).shape == dims
 
     def test_softmax_f16(self):
         x = T.matrix('x', 'float16')