Update tensor.where to allow for case with only condition

pytensor/tensor/basic.py

@@ -760,7 +760,16 @@ def switch(cond, ift, iff):
     """if cond then ift else iff"""
 
 
-where = switch
+def where(cond, ift=None, iff=None):
+    # Normal switch incase both arguements are passed
+    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

@@ -87,6 +87,7 @@
     triu_indices,
     triu_indices_from,
     vertical_stack,
+    where,
     zeros_like,
 )
 from pytensor.tensor.blockwise import Blockwise
@@ -4608,3 +4609,23 @@ def core_np(x, y):
         vectorize_pt(x_test, y_test),
         vectorize_np(x_test, y_test),
     )
+
+
+@pytest.mark.parametrize(
+    "ift, iff",
+    [(None, None), (7, 10), (7, None)],
+    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)