Update tensor.where to allow for case with only condition

pytensor/tensor/basic.py

@@ -760,7 +760,16 @@
     """if cond then ift else iff"""
 
 
-where = switch
+def where(cond, ift=None, iff=None):
+    # Raise an error if only one of ift or iff is passed
+    if (ift is None and iff is not None) or (ift is not None and iff is None):
+        raise Exception("Either both or none of the parameters should be passed")
+    # Normal switch incase both arguements are passed
+    elif ift is not None and iff is not None:
+        return switch(cond, ift, iff)
+    # Add case when only condition is passed
+    else:
+        return as_tensor(cond).nonzero()
 
 
 @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,11 @@ def core_np(x, y):
         vectorize_pt(x_test, y_test),
         vectorize_np(x_test, y_test),
     )
+
+
+def test_where_for_only_condition():
+    a = np.array([1, 2, 3, 4, 5])
+    cond = a <= 3
+    pt_result = where(cond)[0].eval()
+    np_result = np.where(cond)[0]
+    np.testing.assert_allclose(pt_result, np_result)