Enable functional torch.where. (#4298)

docs/source/torch.rst

@@ -40,6 +40,7 @@ Indexing, Slicing, Joining, Mutating Ops
 .. autofunction:: transpose
 .. autofunction:: unbind
 .. autofunction:: unsqueeze
+.. autofunction:: where
 
 
 Random sampling

test/test_autograd.py

@@ -1734,6 +1734,26 @@ def as_strided(x):
         gradcheck(as_strided, [x], raise_exception=True)
         gradgradcheck(as_strided, [x], [Variable(torch.randn(3, 3))])
 
+    def _test_where_functional(self, t):
+        x = Variable(t(torch.randn(5, 5)), requires_grad=True)
+        y = Variable(t(torch.randn(5, 5)), requires_grad=True)
+        cond = Variable(t(mask_not_all_zeros((5, 5))), requires_grad=False)
+
+        def where(cond, x, y):
+            return torch.where(cond, x, y)
+
+        gradcheck(where, [cond, x, y], raise_exception=True)
+        gradgradcheck(where, [cond, x, y], [Variable(t(torch.randn(5, 5)))])
+
+        x = Variable(t(torch.randn(5, 1, 5)), requires_grad=True)
+        y = Variable(t(torch.randn(5, 5, 1)), requires_grad=True)
+        gradcheck(where, [cond, x, y], raise_exception=True)
+        gradgradcheck(where, [cond, x, y], [Variable(t(torch.randn(5, 5, 5)))])
+
+    def test_where_functional(self):
+        # TODO: .cuda() lambda
+        self._test_where_functional(lambda t: t)
+
     def test_inplace_view_backprop_base(self):
         # modify view and back-prop through base
         root = Variable(torch.randn(2, 2), requires_grad=True)
@@ -2398,6 +2418,7 @@ def maybe_non_contig(tensor):
     'baddbmm',
     'addmv',
     'addr',
+    'where'  # argument order
 }
 EXCLUDE_GRADCHECK = {
 }

torch/functional.py

@@ -6,7 +6,7 @@
 
 __all__ = [
     'split', 'chunk', 'stack', 'unbind', 'btriunpack', 'matmul', 'det', 'stft',
-    'hann_window', 'hamming_window', 'bartlett_window',
+    'hann_window', 'hamming_window', 'bartlett_window', 'where',
 ]
 
 
@@ -440,3 +440,30 @@ def bartlett_window(window_length, periodic=True):
         return window[:-1]
     else:
         return window
+
+
+def where(condition, x, y):
+    r"""Return a tensor of elements selected from either :attr:`x` or :attr:`y`, depending on :attr:`condition`.
+
+    defined as::
+
+         out_i =  x_i        if condition_i
+                  y_i        else
+
+    .. note::
+        This function only works with ``Variables``.
+
+    .. note::
+        The tensors :attr:`condition`, :attr:`x`, :attr:`y` must be :ref:`broadcastable <broadcasting-semantics>`.
+
+    Arguments:
+        condition (ByteTensor): When True (nonzero), yield x, otherwise yield y.
+        x (Tensor): values selected at indices where :attr:`condition` is True.
+        y (Tensor): values selected at indices where :attr:`condition` is False.
+
+    Returns:
+        Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`x`, :attr:`y`
+    """
+    # the parameter order is changed here; the functional order is the same as numpy; the
+    # method follows the usual torch mask semantics of x.fn(mask, y)
+    return torch._C._VariableBase.where(x, condition, y)