Repository: chainer/chainer

Issue:
axis (and/or axes) specification for softmax
Currently `functions.softmax` only calculates [w.r.t. the second axis](http://docs.chainer.org/en/latest/_modules/chainer/functions/activation/softmax.html),
but I came across a situation where I need the calculation along other axes (specifically, spatial softmax along the pixel direction).
Although swapping axes does work as a workaround, it would be useful to support an optional argument, like so:

```py
F.softmax(x, axes=(2, 3))
```

Just FYI, Tensorflow [has](https://www.tensorflow.org/api_docs/python/tf/nn/softmax)  an option for such requirements.


Discussion:

Resolving patch:
diff --git a/chainer/functions/activation/softmax.py b/chainer/functions/activation/softmax.py
--- a/chainer/functions/activation/softmax.py
+++ b/chainer/functions/activation/softmax.py
@@ -16,8 +16,9 @@ class Softmax(function.Function):
 
     """Softmax activation function."""
 
-    def __init__(self, use_cudnn=True):
+    def __init__(self, use_cudnn=True, axis=1):
         self.use_cudnn = use_cudnn
+        self.axis = axis
 
     def check_type_forward(self, in_types):
         type_check.expect(in_types.size() == 1)
@@ -26,6 +27,7 @@ def check_type_forward(self, in_types):
         type_check.expect(
             x_type.dtype.kind == 'f',
             x_type.ndim > 1,
+            self.axis < x_type.ndim
         )
 
     def forward(self, x):
@@ -36,17 +38,17 @@ def forward(self, x):
             one = numpy.array(1, dtype=oz_dtype).ctypes
             zero = numpy.array(0, dtype=oz_dtype).ctypes
             handle = cudnn.get_handle()
-            x_cube = x[0].reshape(x[0].shape[:2] + (-1, 1))
-            desc = cudnn.create_tensor_descriptor(x_cube)
+            x_tensor4d = x[0].reshape(self._get_tensor4d_shape(x[0].shape))
+            desc = cudnn.create_tensor_descriptor(x_tensor4d)
             self.y = xp.empty_like(x[0])
             libcudnn.softmaxForward(
                 handle, _algorithm, _mode, one.data, desc.value,
-                x_cube.data.ptr, zero.data, desc.value,
+                x_tensor4d.data.ptr, zero.data, desc.value,
                 self.y.data.ptr)
         else:
-            self.y = x[0] - x[0].max(axis=1, keepdims=True)
+            self.y = x[0] - x[0].max(axis=self.axis, keepdims=True)
             xp.exp(self.y, out=self.y)
-            self.y /= self.y.sum(axis=1, keepdims=True)
+            self.y /= self.y.sum(axis=self.axis, keepdims=True)
 
         return self.y,
 
@@ -59,24 +61,31 @@ def backward(self, x, gy):
             zero = numpy.array(0, dtype=oz_dtype).ctypes
             handle = cudnn.get_handle()
             gx = xp.empty_like(x[0])
-            gx_cube = gx.reshape(gx.shape[:2] + (-1, 1))
-            desc = cudnn.create_tensor_descriptor(gx_cube)
+            gx_tensor4d = gx.reshape(self._get_tensor4d_shape(gx.shape))
+            desc = cudnn.create_tensor_descriptor(gx_tensor4d)
             libcudnn.softmaxBackward(
                 handle, _algorithm, _mode, one.data, desc.value,
                 self.y.data.ptr, desc.value, gy[0].data.ptr, zero.data,
                 desc.value, gx.data.ptr)
         else:
             gx = self.y * gy[0]
-            sumdx = gx.sum(axis=1, keepdims=True)
+            sumdx = gx.sum(axis=self.axis, keepdims=True)
             gx -= self.y * sumdx
 
         return gx,
 
+    def _get_tensor4d_shape(self, shape):
+        left_shape = numpy.prod(shape[slice(0, self.axis)], dtype=numpy.int)
+        center_shape = shape[self.axis]
+        right_shape = numpy.prod(
+            shape[slice(self.axis + 1, len(shape))], dtype=numpy.int)
+        return left_shape, center_shape, right_shape, 1
 
-def softmax(x, use_cudnn=True):
-    """Channelwise softmax function.
 
-    This function computes its softmax along the second axis. Let
+def softmax(x, use_cudnn=True, axis=1):
+    """Softmax function.
+
+    This function computes its softmax along an axis. Let
     :math:`x = (x_1, x_2, \\dots, x_d)^{\\top}` be the d dimensional index
     array and :math:`f(x)` be the d dimensional input array. For each index
     :math:`x` of the input array :math:`f(x)`, it computes the probability
@@ -87,9 +96,10 @@ def softmax(x, use_cudnn=True):
         x (~chainer.Variable): Input variable.
         use_cudnn (bool): If ``True`` and cuDNN is enabled, then this function
             uses cuDNN as the core implementation.
+        axis: The axis along which the softmax is to be computed.
 
     Returns:
         ~chainer.Variable: Output variable.
 
     """
-    return Softmax(use_cudnn)(x)
+    return Softmax(use_cudnn=use_cudnn, axis=axis)(x)


Test patch:
diff --git a/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py b/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py
--- a/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py
+++ b/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py
@@ -13,12 +13,18 @@
 
 
 @testing.parameterize(*testing.product({
-    'shape': [None, (2, 3), (2, 3, 4), (2, 3, 4, 5)],
+    'shape_axis':
+        [{'shape': None, 'axis': 1}, ] +
+        testing.product({'shape': ((3, 4),), 'axis': (0, 1)}) +
+        testing.product({'shape': ((3, 4, 5),), 'axis': (0, 1, 2)}) +
+        testing.product({'shape': ((3, 4, 5, 6),), 'axis': (0, 1, 2, 3)}),
     'dtype': [numpy.float16, numpy.float32, numpy.float64],
 }))
 class TestSoftmax(unittest.TestCase):
 
     def setUp(self):
+        self.shape = self.shape_axis['shape']
+        self.axis = self.shape_axis['axis']
         if self.shape is None:
             # For checking numerical stability
             value = -5 if self.dtype == numpy.float16 else -1000
@@ -36,11 +42,11 @@ def setUp(self):
 
     def check_forward(self, x_data, use_cudnn=True):
         x = chainer.Variable(x_data)
-        y = functions.softmax(x, use_cudnn)
+        y = functions.softmax(x, use_cudnn=use_cudnn, axis=self.axis)
         self.assertEqual(y.data.dtype, self.dtype)
 
         y_expect = numpy.exp(self.x)
-        y_roll = numpy.rollaxis(y_expect, 1, y_expect.ndim)
+        y_roll = numpy.rollaxis(y_expect, self.axis, y_expect.ndim)
         for i in numpy.ndindex(y_roll.shape[:-1]):
             y_roll[i] /= y_roll[i].sum()
 
@@ -63,7 +69,8 @@ def test_forward_gpu_no_cudnn(self):
 
     def check_backward(self, x_data, gy_data, use_cudnn=True):
         gradient_check.check_backward(
-            functions.Softmax(use_cudnn), x_data, gy_data,
+            functions.Softmax(use_cudnn=use_cudnn,
+                              axis=self.axis), x_data, gy_data,
             **self.check_backward_options)
 
     @condition.retry(10)
@@ -82,6 +89,7 @@ def test_backward_gpu_no_cudnn(self):
 
 
 @testing.parameterize(*testing.product({
+    'axis': [0, 1],
     'use_cudnn': [True, False],
     'dtype': [numpy.float16, numpy.float32, numpy.float64],
 }))
@@ -97,7 +105,7 @@ def setUp(self):
 
     def forward(self):
         x = chainer.Variable(self.x)
-        return functions.softmax(x, use_cudnn=self.use_cudnn)
+        return functions.softmax(x, use_cudnn=self.use_cudnn, axis=self.axis)
 
     def test_call_cudnn_forward(self):
         with mock.patch('cupy.cudnn.cudnn.softmaxForward') as func: