Implement torch.broadcast_tensors (#10075)

zou3519 · facebook-github-bot · commit 6b338c802642 · 2018-08-01T19:18:34.000-07:00
Summary: This exposes expand_outplace to python. Fixes #8076. Fixes #10041. I didn't name it torch.broadcast because numpy.broadcast does something slightly different (it returns an object with the correct shape information). Pull Request resolved: #10075 Differential Revision: D9125816 Pulled By: zou3519 fbshipit-source-id: ebe17c8bb54a73ec84b8f76ce14aff3e9c56f4d1
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -12,6 +12,10 @@
 namespace at {
 namespace native {
 
+std::vector<Tensor> broadcast_tensors(TensorList tensors) {
+  return expand_outplace(tensors);
+}
+
 static void check_cat_no_zero_dim(TensorList tensors) {
   for(size_t i = 0; i < tensors.size(); ++i) {
     auto& t = tensors[i];
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -249,6 +249,9 @@
 - func: blackman_window(int64_t window_length, bool periodic, TensorOptions options={}) -> Tensor
   variants: function
 
+- func: broadcast_tensors(TensorList tensors) -> TensorList
+  variants: function
+
 - func: cat(TensorList tensors, int64_t dim=0) -> Tensor
   variants: function
 
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -1850,6 +1850,16 @@ def backward(ctx, grad_output):
         out.sum().backward()
         self.assertEqual(x.grad.data, y_data)
 
+    def test_broadcast_tensors(self):
+        f_args_variable = (torch.randn(3, requires_grad=True),
+                           torch.randn(1, 2, 1, requires_grad=True),
+                           torch.randn(1, 1, requires_grad=True),
+                           torch.randn(5, 1, 1, requires_grad=True))
+        f_args_tensor = deepcopy(unpack_variables(f_args_variable))
+        run_functional_checks(self, "test_broadcast_tensors", "broadcast",
+                              lambda a, b, c, d: torch.broadcast_tensors(a, b, c, d),
+                              True, f_args_variable, f_args_tensor)
+
     def test_cat(self):
         f_args_variable = (torch.randn(1, S, S, requires_grad=True),
                            torch.randn(2, S, S, requires_grad=True),
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -2939,6 +2939,17 @@ def test_broadcast_empty(self):
                          torch.randn(0, 7, 0, 6, 5, 0, 1) + torch.randn(1, 1, 5, 1, 7))
         self.assertRaises(RuntimeError, lambda: torch.randn(7, 0) + torch.randn(2, 1))
 
+    def test_broadcast_tensors(self):
+        x0 = torch.randn(2, 1, 3)
+        x1 = torch.randn(3)
+        x2 = torch.randn(3, 1)
+        expected_size = (2, 3, 3)
+
+        y0, y1, y2 = torch.broadcast_tensors(x0, x1, x2)
+        self.assertTrue(y0.size() == expected_size)
+        self.assertTrue(y1.size() == expected_size)
+        self.assertTrue(y2.size() == expected_size)
+
     @staticmethod
     def _test_contiguous(self, cast):
         x = cast(torch.randn(1, 16, 5, 5))
diff --git a/torch/distributions/utils.py b/torch/distributions/utils.py
@@ -32,30 +32,19 @@ def _finfo(tensor):
     return _FINFO[tensor.storage_type()]
 
 
-def _broadcast_shape(shapes):
-    r"""
-    Given a list of tensor sizes, returns the size of the resulting broadcasted
-    tensor.
-
-    Args:
-        shapes (list of torch.Size): list of tensor sizes
-    """
-    shape = torch.Size()
-    for s in shapes:
-        shape = torch._C._infer_size(s, shape)
-    return shape
+# promote numbers to tensors of dtype torch.get_default_dtype()
+def _default_promotion(v):
+    return torch.tensor(v, dtype=torch.get_default_dtype())
 
 
 def broadcast_all(*values):
     r"""
     Given a list of values (possibly containing numbers), returns a list where each
     value is broadcasted based on the following rules:
-      - `torch.*Tensor` instances are broadcasted as per the `broadcasting rules
-        <http://pytorch.org/docs/master/notes/broadcasting.html>`_
+      - `torch.*Tensor` instances are broadcasted as per :ref:`_broadcasting-semantics`.
       - numbers.Number instances (scalars) are upcast to tensors having
         the same size and type as the first tensor passed to `values`.  If all the
-        values are scalars, then they are upcasted to Tensors having size
-        `(1,)`.
+        values are scalars, then they are upcasted to scalar Tensors.
 
     Args:
         values (list of `numbers.Number` or `torch.*Tensor`)
@@ -64,22 +53,16 @@ def broadcast_all(*values):
         ValueError: if any of the values is not a `numbers.Number` or
             `torch.*Tensor` instance
     """
-    values = list(values)
-    scalar_idxs = [i for i in range(len(values)) if isinstance(values[i], Number)]
-    tensor_idxs = [i for i in range(len(values)) if values[i].__class__.__name__ == 'Tensor']
-    if len(scalar_idxs) + len(tensor_idxs) != len(values):
+    if not all(torch.is_tensor(v) or isinstance(v, Number) for v in values):
         raise ValueError('Input arguments must all be instances of numbers.Number or torch.tensor.')
-    if tensor_idxs:
-        broadcast_shape = _broadcast_shape([values[i].size() for i in tensor_idxs])
-        for idx in tensor_idxs:
-            values[idx] = values[idx].expand(broadcast_shape)
-        template = values[tensor_idxs[0]]
-        for idx in scalar_idxs:
-            values[idx] = template.new(template.size()).fill_(values[idx])
-    else:
-        for idx in scalar_idxs:
-            values[idx] = torch.tensor(float(values[idx]))
-    return values
+    if not all(map(torch.is_tensor, values)):
+        new_tensor = _default_promotion
+        for value in values:
+            if torch.is_tensor(value):
+                new_tensor = value.new_tensor
+                break
+        values = [v if torch.is_tensor(v) else new_tensor(v) for v in values]
+    return torch.broadcast_tensors(*values)
 
 
 def _sum_rightmost(value, dim):
diff --git a/torch/functional.py b/torch/functional.py
@@ -10,6 +10,7 @@
     'argmin',
     'btrifact',
     'btriunpack',
+    'broadcast_tensors',
     'isfinite',
     'isinf',
     'isnan',
@@ -19,6 +20,28 @@
 ]
 
 
+def broadcast_tensors(*tensors):
+    r"""broadcast_tensors(*tensors) -> List of Tensors
+
+    Broadcasts the given tensors according to :ref:`_broadcasting-semantics`.
+
+    Args:
+        *tensors: any number of tensors of the same type
+
+    Example::
+
+        >>> x = torch.arange(3).view(1, 3)
+        >>> y = torch.arange(2).view(2, 1)
+        >>> a, b = torch.broadcast_tensors(x, y)
+        >>> a.size()
+        torch.Size([2, 3])
+        >>> a
+        tensor([[0, 1, 2],
+                [0, 1, 2]])
+    """
+    return torch._C._VariableFunctions.broadcast_tensors(tensors)
+
+
 def split(tensor, split_size_or_sections, dim=0):
     r"""Splits the tensor into chunks.
 
