test for switch with mlx

Author: williambdean

tests/link/mlx/test_basic.py

@@ -5,14 +5,16 @@
 import pytest
 
 from pytensor.compile.function import function
-from pytensor.compile.mode import Mode
+from pytensor.compile.mode import MLX, Mode
+from pytensor.graph import RewriteDatabaseQuery
 from pytensor.graph.basic import Variable
 from pytensor.link.mlx import MLXLinker
 
 
 mx = pytest.importorskip("mlx.core")
 
-mlx_mode = Mode(linker=MLXLinker())
+optimizer = RewriteDatabaseQuery(include=["mlx"], exclude=MLX._optimizer.exclude)
+mlx_mode = Mode(linker=MLXLinker(), optimizer=optimizer)
 py_mode = Mode(linker="py", optimizer=None)
 
 

tests/link/mlx/test_math.py

@@ -42,6 +42,18 @@ def test_elemwise_one_input(op) -> None:
     compare_mlx_and_py([x], out, [x_test])
 
 
+def test_switch() -> None:
+    x = pt.vector("x")
+    y = pt.vector("y")
+
+    out = pt.switch(x > 0, y, x)
+
+    x_test = mx.array([-1.0, 2.0, 3.0])
+    y_test = mx.array([4.0, 5.0, 6.0])
+
+    compare_mlx_and_py([x, y], out, [x_test, y_test])
+
+
 @pytest.mark.parametrize(
     "op",
     [