From 30bd5e5a35e8def113027ec5fc02126f1f68a24b Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 24 Aug 2026 08:57:22 -0500 Subject: [PATCH 1/4] Test shared variables and updates in the MLX backend --- tests/link/mlx/test_basic.py | 41 +++++++++++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/link/mlx/test_basic.py b/tests/link/mlx/test_basic.py index 0aa55c54eb..6f10188abb 100644 --- a/tests/link/mlx/test_basic.py +++ b/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,42 @@ 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 From 00f16ea5bc16d2495155943dcf5b222e69bf970d Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 24 Aug 2026 08:57:40 -0500 Subject: [PATCH 2/4] Convert MLX shared variable updates back to numpy arrays The value written into a shared variable's container outlives the call and may later be read by a function compiled for another backend, so JITLinker.output_filter is applied to update outputs while returned outputs keep their native backend type. --- pytensor/link/basic.py | 22 +++++++++++++++++++++- pytensor/link/mlx/linker.py | 7 +++++++ tests/link/mlx/test_basic.py | 24 ++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/pytensor/link/basic.py b/pytensor/link/basic.py index 80e57afff1..0b4563fb43 100644 --- a/pytensor/link/basic.py +++ b/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. + + Values returned to the caller keep whatever native type the backend produced, + but a value written back into a shared variable's container outlives the call + and may later be read by a function compiled for another backend. Backends whose + arrays are not `numpy` ones override this; the default passes the value through. + """ return out def create_jitable_thunk( @@ -665,12 +671,22 @@ 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. + filters_output = type(self).output_filter is not JITLinker.output_filter + update_output_idxs = ( + tuple(self.fgraph.update_mapping or ()) if filters_output 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 +699,10 @@ def thunk( for o_storage, o_val in zip(thunk_outputs, outputs): o_storage[0] = o_val + for idx in update_output_idxs: + o_storage = thunk_outputs[idx] + o_storage[0] = output_filter(fgraph_outputs[idx], o_storage[0]) + else: # Edge case - functions without outputs def thunk( diff --git a/pytensor/link/mlx/linker.py b/pytensor/link/mlx/linker.py index 9d662308cf..70211bd18e 100644 --- a/pytensor/link/mlx/linker.py +++ b/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. diff --git a/tests/link/mlx/test_basic.py b/tests/link/mlx/test_basic.py index 6f10188abb..e4ce5bbd28 100644 --- a/tests/link/mlx/test_basic.py +++ b/tests/link/mlx/test_basic.py @@ -394,3 +394,27 @@ def test_shared_updates(): 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.container.storage[0], np.ndarray) + assert isinstance(a.get_value(borrow=True), np.ndarray) + np.testing.assert_allclose(a.get_value(), np.array([2, 3, 4], 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)) From 95e4f9d57ce356f15417ca9f2f153b73acf03493 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 24 Aug 2026 18:47:13 -0500 Subject: [PATCH 3/4] Clarify the JITLinker.output_filter contract --- pytensor/link/basic.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/pytensor/link/basic.py b/pytensor/link/basic.py index 0b4563fb43..fd5b00b6d9 100644 --- a/pytensor/link/basic.py +++ b/pytensor/link/basic.py @@ -619,12 +619,12 @@ def input_filter(self, inp: Any) -> Any: return inp def output_filter(self, var: Variable, out: Any) -> Any: - """Convert a value the JITed function produced into one `pytensor` can store. + """Convert a value the JITed function produced into one PyTensor can store. - Values returned to the caller keep whatever native type the backend produced, - but a value written back into a shared variable's container outlives the call - and may later be read by a function compiled for another backend. Backends whose - arrays are not `numpy` ones override this; the default passes the value through. + 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 @@ -673,9 +673,11 @@ def create_jitable_thunk( # Shared variable updates are the only outputs worth converting, and a backend # that leaves `output_filter` alone pays nothing for the hook. - filters_output = type(self).output_filter is not JITLinker.output_filter + overrides_output_filter = ( + type(self).output_filter is not JITLinker.output_filter + ) update_output_idxs = ( - tuple(self.fgraph.update_mapping or ()) if filters_output else () + tuple(self.fgraph.update_mapping or ()) if overrides_output_filter else () ) if thunk_outputs: @@ -700,8 +702,10 @@ def thunk( o_storage[0] = o_val for idx in update_output_idxs: - o_storage = thunk_outputs[idx] - o_storage[0] = output_filter(fgraph_outputs[idx], o_storage[0]) + update_storage = thunk_outputs[idx] + update_storage[0] = output_filter( + fgraph_outputs[idx], update_storage[0] + ) else: # Edge case - functions without outputs From 29fc4e41c757b2569c5f65ed26d1a149972fbd39 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 24 Aug 2026 18:47:18 -0500 Subject: [PATCH 4/4] Cover multiple shared variable updates in the MLX backend --- tests/link/mlx/test_basic.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/link/mlx/test_basic.py b/tests/link/mlx/test_basic.py index e4ce5bbd28..fcda129e78 100644 --- a/tests/link/mlx/test_basic.py +++ b/tests/link/mlx/test_basic.py @@ -405,11 +405,27 @@ def test_shared_updates_are_not_device_arrays(): # 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.container.storage[0], np.ndarray) 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))