Repository: chainer/chainer

Issue:
Make the number of axes for batch dimensions configurable in Linear
L.Linear accepts multiple channel dimensions. It is useful in computer vision applications where the output of the final convolution layer can be directly fed to a linear link. This feature can be seen as a flexible number of channel axes.

On the other hand, the number of batch axes is currently fixed to one, which would be cumbersome in some use cases, especially when we have to deal with a sequence of input arrays. In this case, the first two axes may be used to represent the mini-batch and the sequence. We currently need to reshape it before and after feeding it to the linear link, which has overheads both for coding and runtime.

It can be resolved by introducing an option to Linear (and also to the F.linear function) that indicates the number of batch axes.

The same discussions would also be applied to convolution links/functions.


Discussion:
Shouldn't the same discussions also be applied to all the other links and functions that assume the first axis is the only batch axis e.g. F.huber_loss and F.select_item?
You're right. It should be applied to other functions that assume the batch axis as well.
This issue has been automatically marked as stale because it has not had recent activity. It will be closed after 30 days if no further activity occurs. Thank you for your contributions.


Resolving patch:
diff --git a/chainer/functions/connection/linear.py b/chainer/functions/connection/linear.py
--- a/chainer/functions/connection/linear.py
+++ b/chainer/functions/connection/linear.py
@@ -185,7 +185,7 @@ def backward(self, indexes, grad_outputs):
         return ret
 
 
-def linear(x, W, b=None):
+def linear(x, W, b=None, n_batch_axes=1):
     """Linear function, or affine transformation.
 
     It accepts two or three arguments: an input minibatch ``x``, a weight
@@ -195,21 +195,25 @@ def linear(x, W, b=None):
 
     Args:
         x (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
-        :class:`cupy.ndarray`): Input variable, which is a :math:`(s_B, s_1, \
-            s_2, ..., s_n)`-shaped float array. Its first dimension
-            :math:`(s_B)` is assumed to be the *minibatch dimension*. The
-            other dimensions are treated as concatenated one dimension whose
-            size must be :math:`(s_1 * ... * s_n = N)`.
+        :class:`cupy.ndarray`): Input variable, which is a :math:`(s_1, s_2, \
+            ..., s_n)`-shaped float array. Its first ``n_batch_axes``
+            dimensions are handled as *minibatch dimensions*. The
+            other dimensions are handled as concatenated one dimension whose
+            size must be :math:`(s_{\\rm n_batch_axes} * ... * s_n = N)`.
         W (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
         :class:`cupy.ndarray`): Weight variable of shape :math:`(M, N)`,
-            where :math:`(N = s_1 * ... * s_n)`.
+            where :math:`(N = s_{\\rm n_batch_axes} * ... * s_n)`.
         b (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
         :class:`cupy.ndarray`): Bias variable (optional) of shape
             :math:`(M,)`.
+        n_batch_axes (int): The number of batch axes. The default is 1. The
+            input variable is reshaped into
+            :math:`{\\rm n_batch_axes} + 1`-dimensional tensor.
+            This should be greater than 0.
 
     Returns:
         ~chainer.Variable: Output variable. A float array with shape
-        of :math:`(s_B, M)`.
+        of :math:`(s_1, ..., s_{\\rm n_batch_axes}, M)`.
 
     .. seealso:: :class:`~chainer.links.Linear`
 
@@ -223,13 +227,20 @@ def linear(x, W, b=None):
         (3, 5)
 
     """
-    if x.ndim > 2:
-        x = x.reshape(len(x), -1)
-
+    if n_batch_axes <= 0:
+        raise ValueError('n_batch_axes should be greater than 0.')
+    if n_batch_axes > 1:
+        batch_shape = x.shape[:n_batch_axes]
+        batch_size = numpy.prod(batch_shape)
+        x = x.reshape(batch_size, -1)
+    elif x.ndim > 2:
+        x = x.reshape(x.shape[0], -1)
     if b is None:
         args = x, W
     else:
         args = x, W, b
 
     y, = LinearFunction().apply(args)
+    if n_batch_axes > 1:
+        y = y.reshape(batch_shape + (-1,))
     return y


Test patch:
diff --git a/tests/chainer_tests/functions_tests/connection_tests/test_linear.py b/tests/chainer_tests/functions_tests/connection_tests/test_linear.py
--- a/tests/chainer_tests/functions_tests/connection_tests/test_linear.py
+++ b/tests/chainer_tests/functions_tests/connection_tests/test_linear.py
@@ -18,6 +18,8 @@ def _to_noncontiguous(arrays):
 @testing.parameterize(*testing.product({
     'x_dtype': [numpy.float16, numpy.float32, numpy.float64],
     'W_dtype': [numpy.float16, numpy.float32, numpy.float64],
+    'x_shape': [{'n_batch_axes': 1, 'data_shape': (3,)},
+                {'n_batch_axes': 3, 'data_shape': (3, 5)}],
     'c_contiguous': [True, False],
     'nobias': [True, False],
 }))
@@ -35,15 +37,20 @@ def _to_noncontiguous(arrays):
 class TestNonparameterizedLinear(unittest.TestCase):
 
     def setUp(self):
-        W = numpy.random.uniform(
-            -1, 1, (2, 3)).astype(self.W_dtype)
+        self.n_batch_axes = self.x_shape['n_batch_axes']
+        data_shape = self.x_shape['data_shape']
+        input_size = numpy.prod(data_shape)
+        W = numpy.random.uniform(-1, 1, (2, input_size)).astype(self.W_dtype)
         if self.nobias:
             b = None
         else:
             b = numpy.random.uniform(-1, 1, 2).astype(self.x_dtype)
 
-        x = numpy.random.uniform(-1, 1, (4, 3)).astype(self.x_dtype)
-        gy = numpy.random.uniform(-1, 1, (4, 2)).astype(self.x_dtype)
+        batch_shape = (4,) + (2,) * (self.n_batch_axes - 1)
+        x = numpy.random.uniform(
+            -1, 1, batch_shape + data_shape).astype(self.x_dtype)
+        gy = numpy.random.uniform(
+            -1, 1, batch_shape + (2,)).astype(self.x_dtype)
         ggx = numpy.random.uniform(-1, 1, x.shape).astype(self.x_dtype)
         ggW = numpy.random.uniform(-1, 1, W.shape).astype(self.W_dtype)
         if self.nobias:
@@ -68,24 +75,30 @@ def setUp(self):
         else:
             self.check_forward_options = {}
             self.check_backward_options = {
-                'atol': 1e-3, 'rtol': 1e-3}
+                'atol': 1e-2, 'rtol': 1e-2}
             self.check_double_backward_options = {
-                'atol': 1e-3, 'rtol': 1e-3}
+                'atol': 1e-2, 'rtol': 1e-2}
 
     def forward_cpu(self, inputs):
         x, W, b = inputs
+        if self.n_batch_axes > 1:
+            batch_shape = x.shape[:self.n_batch_axes]
+            batch_size = numpy.prod(batch_shape)
+            x = x.reshape(batch_size, -1)
         y = x.dot(W.T)
         if b is not None:
             y += b
+        if self.n_batch_axes > 1:
+            y = y.reshape(batch_shape + (-1,))
         return y,
 
     def forward(self, *inputs):
         if len(inputs) == 3:
             x, W, b = inputs
-            y = functions.linear(x, W, b)
+            y = functions.linear(x, W, b, n_batch_axes=self.n_batch_axes)
         else:
             x, W = inputs
-            y = functions.linear(x, W)
+            y = functions.linear(x, W, n_batch_axes=self.n_batch_axes)
         return y,
 
     def check_forward(self, inputs, backend_config):
@@ -182,4 +195,23 @@ def test_1(self):
             z.backward()
 
 
+class TestLinearNBatchAxesBoundaryCondition(unittest.TestCase):
+
+    def setUp(self):
+        self.W = numpy.random.uniform(
+            -1, 1, (2, 15)).astype(numpy.float32)
+        self.x = numpy.random.uniform(
+            -1, 1, (3, 3, 5)).astype(numpy.float32)
+
+    def test_negative(self):
+        n_batch_axes = -1
+        with self.assertRaises(ValueError):
+            functions.linear(self.x, self.W, n_batch_axes=n_batch_axes)
+
+    def test_zero(self):
+        n_batch_axes = 0
+        with self.assertRaises(ValueError):
+            functions.linear(self.x, self.W, n_batch_axes=n_batch_axes)
+
+
 testing.run_module(__name__, __file__)