Repository: pytorch/pytorch

Issue:
I think Bilinear() should support 3D tensor like Linear() does
When I execute following code, it is gonna get 

> RuntimeError: matrices expected, got 3D, 2D tensors

`import torch.nn as nn`
`m = nn.Bilinear(2, 3, 4)`
`input1 = Variable(torch.randn(5, 3, 2))`
`input2 = Variable(torch.randn(5, 3, 3))`
`output = m(input1, input2)`
`print(output.size())`
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-23-b247cd5ad730> in <module>()
----> 1 output = m(input1, input2)
      2 print(output.size())

~/anaconda3/envs/deepchem/lib/python3.5/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    323         for hook in self._forward_pre_hooks.values():
    324             hook(self, input)
--> 325         result = self.forward(*input, **kwargs)
    326         for hook in self._forward_hooks.values():
    327             hook_result = hook(self, input, result)

~/anaconda3/envs/deepchem/lib/python3.5/site-packages/torch/nn/modules/linear.py in forward(self, input1, input2)
    110 
    111     def forward(self, input1, input2):
--> 112         return F.bilinear(input1, input2, self.weight, self.bias)
    113 
    114     def __repr__(self):

~/anaconda3/envs/deepchem/lib/python3.5/site-packages/torch/nn/functional.py in bilinear(input1, input2, weight, bias)
    845         return Bilinear.apply(input1, input2, weight)
    846     else:
--> 847         return Bilinear.apply(input1, input2, weight, bias)
    848 
    849 

~/anaconda3/envs/deepchem/lib/python3.5/site-packages/torch/nn/_functions/linear.py in forward(ctx, input1, input2, weight, bias)
     15         # compute output scores:
     16         for k, w in enumerate(weight):
---> 17             torch.mm(input1, w, out=buff)
     18             buff.mul_(input2)
     19             torch.sum(buff, 1, keepdim=True, out=output.narrow(1, k, 1))

RuntimeError: matrices expected, got 3D, 2D tensors at /opt/conda/conda-bld/pytorch_1512383260527/work/torch/lib/TH/generic/THTensorMath.c:1411


Discussion:

Resolving patch:
diff --git a/torch/nn/_functions/linear.py b/torch/nn/_functions/linear.py
deleted file mode 100644
--- a/torch/nn/_functions/linear.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import torch
-from torch.autograd import Function, Variable
-
-
-class Bilinear(Function):
-
-    @staticmethod
-    def forward(ctx, input1, input2, weight, bias=None):
-        ctx.save_for_backward(input1, input2, weight, bias)
-
-        output = input1.new(input1.size(0), weight.size(0))
-
-        buff = input1.new()
-
-        # compute output scores:
-        for k, w in enumerate(weight):
-            torch.mm(input1, w, out=buff)
-            buff.mul_(input2)
-            torch.sum(buff, 1, keepdim=True, out=output.narrow(1, k, 1))
-
-        if bias is not None:
-            output.add_(bias.expand_as(output))
-
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        input1, input2, weight, bias = ctx.saved_variables
-        grad_input1 = grad_input2 = grad_weight = grad_bias = None
-
-        buff = Variable(input1.data.new())
-
-        if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
-            grad_input1 = torch.mm(input2, weight[0].t())
-            grad_input1 = grad_input1.mul(grad_output.narrow(1, 0, 1).expand(grad_input1.size()))
-            grad_input2 = torch.mm(input1, weight[0])
-            grad_input2 = grad_input2.mul(grad_output.narrow(1, 0, 1).expand(grad_input2.size()))
-
-            for k in range(1, weight.size(0)):
-                buff = input2.mm(weight[k].t())
-                buff = buff.mul(grad_output.narrow(1, k, 1).expand(grad_input1.size()))
-                grad_input1.add_(buff)
-
-                buff = input1.mm(weight[k])
-                buff = buff.mul(grad_output.narrow(1, k, 1).expand(grad_input2.size()))
-                grad_input2.add_(buff)
-
-        grad_weight = Variable(weight.data.new(weight.size()))
-        if ctx.needs_input_grad[2]:
-            # accumulate parameter gradients:
-            for k in range(weight.size(0)):
-                buff = input1.mul(grad_output.narrow(1, k, 1).expand_as(input1))
-                grad_weight[k] = torch.mm(buff.t(), input2)
-
-        if bias is not None and ctx.needs_input_grad[3]:
-            grad_bias = grad_output.sum(0, keepdim=False)
-
-        return grad_input1, grad_input2, grad_weight, grad_bias
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
--- a/torch/nn/functional.py
+++ b/torch/nn/functional.py
@@ -9,7 +9,6 @@
 from torch._C import _infer_size, _add_docstr
 from . import _functions
 from .modules import utils
-from ._functions.linear import Bilinear
 from ._functions.padding import ConstantPadNd
 from ._functions import vision
 from ._functions.thnn.fold import Col2Im, Im2Col
@@ -1001,10 +1000,7 @@ def linear(input, weight, bias=None):
 
 
 def bilinear(input1, input2, weight, bias=None):
-    if bias is None:
-        return Bilinear.apply(input1, input2, weight)
-    else:
-        return Bilinear.apply(input1, input2, weight, bias)
+    return torch._C._VariableFunctions.bilinear(input1, input2, weight, bias)
 
 
 def embedding(input, weight, padding_idx=None, max_norm=None, norm_type=2,
diff --git a/torch/nn/modules/linear.py b/torch/nn/modules/linear.py
--- a/torch/nn/modules/linear.py
+++ b/torch/nn/modules/linear.py
@@ -73,8 +73,11 @@ class Bilinear(Module):
             Default: ``True``
 
     Shape:
-        - Input: :math:`(N, \text{in1_features})`, :math:`(N, \text{in2_features})`
-        - Output: :math:`(N, \text{out_features})`
+        - Input: :math:`(N, *, \text{in1_features})`, :math:`(N, *, \text{in2_features})`
+          where :math:`*` means any number of additional dimensions. All but the last
+          dimension of the inputs should be the same.
+        - Output: :math:`(N, *, \text{out_features})` where all but the last dimension
+          are the same shape as the input.
 
     Attributes:
         weight: the learnable weights of the module of shape


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -4145,6 +4145,13 @@ def test_bilinear(self):
         _assertGradAndGradgradChecks(self, lambda x1, x2: F.bilinear(x1, x2, module.weight, module.bias),
                                      (input1_1, input2_1))
 
+    def test_bilinear_broadcasting(self):
+        m = nn.Bilinear(5, 6, 8)
+        input1 = torch.randn(2, 3, 5)
+        input2 = torch.randn(2, 3, 6)
+        expected = m(input1.view(6, 5), input2.view(6, 6)).view(2, 3, 8)
+        self.assertEqual(expected, m(input1, input2))
+
     def test_conv_tbc(self):
         inp = Variable(torch.randn(9, 4, 5), requires_grad=True)
         weight = Variable(torch.randn(3, 5, 6), requires_grad=True)