Support N-D tensors in Bilinear (#5764)

li-roy · soumith · commit 1dcad0853729 · 2018-03-17T11:57:43.000-04:00
* support n-d inputs in bilinear and move to aten

* support n-d inputs in bilinear and move to aten

* add asserts to bilinear inputs

* address comments

* cast int64_t in asserts
diff --git a/aten/src/ATen/native/Linear.cpp b/aten/src/ATen/native/Linear.cpp
@@ -0,0 +1,36 @@
+#include "ATen/ATen.h"
+#include "ATen/NativeFunctions.h"
+
+
+namespace at { namespace native {
+
+Tensor bilinear(const Tensor& input1, const Tensor& input2, const Tensor& weight, const Tensor& bias) {
+  AT_ASSERT(input1.dim() == input2.dim(), "bilinear(): input dimensions do not match: got %lld and %lld",
+            (long long)input1.dim(), (long long)input2.dim());
+  for (int64_t i = 0; i < input1.dim() - 1; i++) {
+    AT_ASSERT(input1.size(i) == input2.size(i),
+              "bilinear(): input batch dimensions do not match at dim %lld: got %lld and %lld",
+              (long long)i, (long long)input1.size(i), (long long)input2.size(i));
+  }
+  AT_ASSERT(input1.size(input1.dim() - 1) == weight.size(1),
+            "bilinear(): input1 size does not match weight size: got %lld but expected %lld",
+            (long long)input1.size(input1.dim() - 1), (long long)weight.size(1));
+  AT_ASSERT(input2.size(input2.dim() - 1) == weight.size(2),
+            "bilinear(): input2 size does not match weight size: got %lld but expected %lld",
+            (long long)input2.size(input2.dim() - 1), (long long)weight.size(2));
+  AT_ASSERT(bias.defined() && bias.size(0) == weight.size(0),
+            "bilinear(): bias size does not match weight size: got %lld but expected %lld",
+            (long long)bias.size(0), (long long)weight.size(0));
+
+  auto b_input1 = input1.unsqueeze(-2).unsqueeze(-2);
+  auto b_input2 = input2.unsqueeze(-2).unsqueeze(-1);
+
+  auto output = at::matmul(at::matmul(b_input1, weight), b_input2);
+  output = output.squeeze(-1).squeeze(-1);
+  if (bias.defined()) {
+    output = output + bias;
+  }
+  return output;
+}
+
+}}  // namespace at::native
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -71,6 +71,9 @@
 
 - func: bernoulli_(Tensor self, double p=0.5, Generator* generator=nullptr) -> Tensor
 
+- func: bilinear(Tensor input1, Tensor input2, Tensor weight, Tensor? bias) -> Tensor
+  variants: function
+
 - func: cat(TensorList tensors, int64_t dim=0) -> Tensor
   variants: function
 
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -4169,6 +4169,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)
diff --git a/torch/nn/_functions/linear.py b/torch/nn/_functions/linear.py
diff --git 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
@@ -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
