Convert MLX shared variable updates back to numpy arrays

pytensor/link/basic.py

@@ -619,7 +619,13 @@ def input_filter(self, inp: Any) -> Any:
         return inp
 
     def output_filter(self, var: Variable, out: Any) -> Any:
-        """Apply a filter to the data output by a JITed function call."""
+        """Convert a value the JITed function produced into one PyTensor can store.
+
+        Only values written back into a shared variable's container pass through here.
+        Those outlive the call and may later be read by a function compiled for another
+        backend, so a backend whose arrays are not NumPy ones overrides this; the default
+        passes the value through.
+        """
         return out
 
     def create_jitable_thunk(
@@ -665,12 +671,24 @@ def create_jitable_thunk(
         thunk_outputs = [storage_map[n] for n in self.fgraph.outputs]
         fgraph_jit = self.jit_compile(converted_fgraph)
 
+        # Shared variable updates are the only outputs worth converting, and a backend
+        # that leaves `output_filter` alone pays nothing for the hook.
+        overrides_output_filter = (
+            type(self).output_filter is not JITLinker.output_filter
+        )
+        update_output_idxs = (
+            tuple(self.fgraph.update_mapping or ()) if overrides_output_filter else ()
+        )
+
         if thunk_outputs:
 
             def thunk(
                 fgraph_jit=fgraph_jit,
                 thunk_inputs=thunk_inputs,
                 thunk_outputs=thunk_outputs,
+                update_output_idxs=update_output_idxs,
+                output_filter=self.output_filter,
+                fgraph_outputs=self.fgraph.outputs,
             ):
                 try:
                     outputs = fgraph_jit(*(x[0] for x in thunk_inputs))
@@ -683,6 +701,12 @@ def thunk(
                 for o_storage, o_val in zip(thunk_outputs, outputs):
                     o_storage[0] = o_val
 
+                for idx in update_output_idxs:
+                    update_storage = thunk_outputs[idx]
+                    update_storage[0] = output_filter(
+                        fgraph_outputs[idx], update_storage[0]
+                    )
+
         else:
             # Edge case - functions without outputs
             def thunk(

pytensor/link/mlx/linker.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from pytensor.link.basic import JITLinker
 
 
@@ -58,6 +60,11 @@ def fn(*inputs, inner_fn=inner_fn):
 
         return fn
 
+    def output_filter(self, var, out):
+        import mlx.core as mx
+
+        return np.asarray(out) if isinstance(out, mx.array) else out
+
     def create_thunk_inputs(self, storage_map):
         """Create inputs for the MLX thunk.
 

tests/link/mlx/test_basic.py

@@ -10,7 +10,7 @@
 import pytest
 
 import pytensor
-from pytensor import config
+from pytensor import config, shared
 from pytensor import tensor as pt
 from pytensor.compile.maker import function
 from pytensor.compile.mode import MLX, Mode
@@ -355,3 +355,82 @@ def test_nan_array_constant():
     compare_mlx_and_py(
         [x], [x + c], [np.array([10.0, 20.0, 30.0], dtype=config.floatX)]
     )
+
+
+def test_shared():
+    a = shared(np.array([1, 2, 3], dtype=config.floatX))
+
+    pytensor_mlx_fn = function([], a, mode=mlx_mode)
+    mlx_res = pytensor_mlx_fn()
+
+    assert isinstance(mlx_res, mx.array)
+    np.testing.assert_allclose(np.asarray(mlx_res), a.get_value())
+
+    pytensor_mlx_fn = function([], a * 2, mode=mlx_mode)
+    mlx_res = pytensor_mlx_fn()
+
+    assert isinstance(mlx_res, mx.array)
+    np.testing.assert_allclose(np.asarray(mlx_res), a.get_value() * 2)
+
+    new_a_value = np.array([3, 4, 5], dtype=config.floatX)
+    a.set_value(new_a_value)
+
+    mlx_res = pytensor_mlx_fn()
+    assert isinstance(mlx_res, mx.array)
+    np.testing.assert_allclose(np.asarray(mlx_res), new_a_value * 2)
+
+
+def test_shared_updates():
+    a = shared(0)
+
+    pytensor_mlx_fn = function([], a, updates={a: a + 1}, mode=mlx_mode)
+    res1, res2 = pytensor_mlx_fn(), pytensor_mlx_fn()
+    assert res1 == 0
+    assert res2 == 1
+    assert a.get_value() == 2
+
+    a.set_value(5)
+    res1, res2 = pytensor_mlx_fn(), pytensor_mlx_fn()
+    assert res1 == 5
+    assert res2 == 6
+    assert a.get_value() == 7
+
+
+def test_shared_updates_are_not_device_arrays():
+    a = shared(np.array([1, 2, 3], dtype=config.floatX))
+
+    pytensor_mlx_fn = function([], a, updates={a: a + 1}, mode=mlx_mode)
+    mlx_res = pytensor_mlx_fn()
+
+    # The returned value stays an MLX array, but the one stored back in the shared
+    # variable must not, or any other consumer of the container chokes on it.
+    assert isinstance(mlx_res, mx.array)
+    assert isinstance(a.get_value(borrow=True), np.ndarray)
+    assert isinstance(a.get_value(borrow=False), np.ndarray)
+    np.testing.assert_allclose(a.get_value(), np.array([2, 3, 4], dtype=config.floatX))
+
+
+def test_multiple_shared_updates():
+    a = shared(np.array([1, 2, 3], dtype=config.floatX))
+    b = shared(np.array([10, 20, 30], dtype=config.floatX))
+    x = pt.vector("x", dtype=config.floatX)
+
+    pytensor_mlx_fn = function([x], x * 2, updates={a: a + 1, b: b * 2}, mode=mlx_mode)
+    mlx_res = pytensor_mlx_fn(np.array([1, 1, 1], dtype=config.floatX))
+
+    assert isinstance(mlx_res, mx.array)
+    for shared_var, expected in ((a, [2, 3, 4]), (b, [20, 40, 60])):
+        assert isinstance(shared_var.get_value(borrow=True), np.ndarray)
+        np.testing.assert_allclose(
+            shared_var.get_value(), np.array(expected, dtype=config.floatX)
+        )
+
+
+def test_shared_updates_readable_by_other_backend():
+    a = shared(np.array([1, 2, 3], dtype=config.floatX))
+
+    mlx_fn = function([], a, updates={a: a + 1}, mode=mlx_mode)
+    other_fn = function([], a * 2)
+
+    mlx_fn()
+    np.testing.assert_allclose(other_fn(), np.array([4, 6, 8], dtype=config.floatX))