Update tensor.where to allow for case with only condition

pytensor/tensor/basic.py

@@ -761,15 +761,30 @@
 
 
 def where(cond, ift=None, iff=None):
-    # Normal switch incase both arguements are passed
+    """
+    where(condition, [ift, iff])
+    Return elements chosen from `ift` or `iff` depending on `condition`.
+
+    Note: When only condition is provided, this function is a shorthand for `as_tensor(condition).nonzero()`.
+
+    Parameters
+    ----------
+    condition : tensor_like, bool
+        Where True, yield `ift`, otherwise yield `iff`.
+    x, y : tensor_like
+        Values from which to choose.
+
+    Returns
+    -------
+    out : TensorVariable
+        A tensor with elements from `ift` where `condition` is True, and elements from `iff` elsewhere.
+    """
     if ift is not None and iff is not None:
         return switch(cond, ift, iff)
-    # Add case when only condition is passed
     elif ift is None and iff is None:
         return as_tensor(cond).nonzero()
-    # Raise an error if only one arguement is passed
     else:
         raise Exception("Either both or none of the parameters should be passed")
 
 
 @scalar_elemwise

tests/tensor/test_basic.py

@@ -4613,19 +4613,12 @@ def core_np(x, y):
 
 @pytest.mark.parametrize(
     "ift, iff",
-    [(None, None), (7, 10), (7, None)],
+    [(None, None), (7, 10), pytest.param(7, None, marks=[pytest.mark.xfail])],
     ids=["both none", "both valid", "one none"],
 )
 def test_where_for_only_condition(ift, iff):
     a = np.array([1, 2, 3, 4, 5])
     cond = a < 3
-    if ift is not None and iff is not None:
-        pt_result = function([], where(cond, ift, iff))()
-        np_result = np.where(cond, ift, iff)
-        np.testing.assert_allclose(pt_result, np_result)
-    elif ift is None and iff is None:
-        pt_result = function([], where(cond))()
-        np_result = np.where(cond)
-        np.testing.assert_allclose(pt_result, np_result)
-    else:
-        pytest.raises(Exception, where, cond, ift)
+    np_result = np.where(*[x for x in [cond, ift, iff] if x is not None])
+    pt_result = function([], where(cond, ift, iff))()
+    np.testing.assert_allclose(np_result, pt_result)