Repository: chainer/chainer

Issue:
[BUG] Wrong initialization of the Linear link weights when called with n_batch_axes>1
# **Overview**
When creating the _Linear_ link without specifying the __input_size__ in the constructor and then calling the link with the parameter __n_batch_axes__ > 1, the weight matrix is initialized with a wrong shape and causes an error.

# **Conditions**
Platform: Windows-7-6.1.7601-SP1
Chainer: 5.4.0
NumPy: 1.15.4
CuPy: Not Available
iDeep: Not Available

# Code to reproduce
```
import numpy as np
import chainer
batch_size, seq_len, num_of_features, out_features = (1, 4, 3, 8)
linear_layer_chainer = chainer.links.Linear(out_features)
data = chainer.Variable(np.ones((batch_size, seq_len, num_of_features), dtype=np.float32))
results = linear_layer_chainer (data, n_batch_axes=2)
```


# **Error messages, stack traces, or logs**
```
Traceback (most recent call last):
  File "C:\Users\ntt\Anaconda3\lib\site-packages\IPython\core\interactiveshell.py", line 3267, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-9-9df2914c2f5a>", line 1, in <module>
    results = linear_layer_chainer_2_axes(data, n_batch_axes=2)
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\link.py", line 242, in __call__
    out = forward(*args, **kwargs)
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\links\connection\linear.py", line 138, in forward
    return linear.linear(x, self.W, self.b, n_batch_axes=n_batch_axes)
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\functions\connection\linear.py", line 288, in linear
    y, = LinearFunction().apply(args)
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\function_node.py", line 245, in apply
    self._check_data_type_forward(in_data)
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\function_node.py", line 330, in _check_data_type_forward
    self.check_type_forward(in_type)
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\functions\connection\linear.py", line 27, in check_type_forward
    x_type.shape[1] == w_type.shape[1],
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\utils\type_check.py", line 546, in expect
    expr.expect()
  File "C:\Users\ntt\Anaconda3\lib\site-packages\chainer\utils\type_check.py", line 483, in expect
    '{0} {1} {2}'.format(left, self.inv, right))
chainer.utils.type_check.InvalidType: 
Invalid operation is performed in: LinearFunction (Forward)
Expect: x.shape[1] == W.shape[1]
Actual: 3 != 12
```

# Others 
If you create the _Linear_ layer link with __in_size__ passed explicitly to the correct shape, it works:

```
working_linear = chainer.links.Linear(num_of_features, out_features)
working_results = working_linear(data,  n_batch_axes=2)
working_results.shape
```
Output : `(1, 4, 8)`

# Solution 
I believe the culprit lays in the _Linear_ link _forward_ method:
```
    def forward(self, x, n_batch_axes=1):
        """Applies the linear layer.

        Args:
            x (~chainer.Variable): Batch of input vectors.
            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 of the linear layer.

        """
        if self.W.array is None:
            in_size = functools.reduce(operator.mul, x.shape[1:], 1)
            self._initialize_params(in_size)
        return linear.linear(x, self.W, self.b, n_batch_axes=n_batch_axes)
```
If I interpret this part correctly, it does not take into account the _n_batch_axes_ parameter when initializing the weight matrix.




Discussion:

Resolving patch:
diff --git a/chainer/links/connection/linear.py b/chainer/links/connection/linear.py
--- a/chainer/links/connection/linear.py
+++ b/chainer/links/connection/linear.py
@@ -150,6 +150,6 @@ def forward(self, x, n_batch_axes=1):
 
         """
         if self.W.array is None:
-            in_size = utils.size_of_shape(x.shape[1:])
+            in_size = utils.size_of_shape(x.shape[n_batch_axes:])
             self._initialize_params(in_size)
         return linear.linear(x, self.W, self.b, n_batch_axes=n_batch_axes)


Test patch:
diff --git a/tests/chainer_tests/links_tests/connection_tests/test_linear.py b/tests/chainer_tests/links_tests/connection_tests/test_linear.py
--- a/tests/chainer_tests/links_tests/connection_tests/test_linear.py
+++ b/tests/chainer_tests/links_tests/connection_tests/test_linear.py
@@ -165,6 +165,17 @@ def test_empty_batch_dim(self):
         assert y.shape == (0, 4)
 
 
+class TestNBatchAxesInitialize(unittest.TestCase):
+
+    def setUp(self):
+        self.link = links.Linear(4)
+        self.x = numpy.random.uniform(-1, 1, (2, 5, 3)).astype(numpy.float32)
+
+    def test_init_n_batch_axes(self):
+        y = self.link(chainer.Variable(self.x), n_batch_axes=2)
+        assert y.shape == (2, 5, 4)
+
+
 class TestInvalidLinear(unittest.TestCase):
 
     def setUp(self):