MLX backend POC #1365

williambdean · 2025-04-11T15:39:43Z

Description

Getting ball rolling started with #1350

Related Issue

Closes #
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pytensor--1365.org.readthedocs.build/en/1365/

pytensor/link/pytorch/linker.py

codecov · 2025-04-11T18:11:07Z

Codecov Report

Attention: Patch coverage is 80.23256% with 17 lines in your changes missing coverage. Please review.

Project coverage is 82.04%. Comparing base (4e59f21) to head (d057453).
Report is 13 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/mlx/linker.py	77.50%	6 Missing and 3 partials ⚠️
pytensor/link/mlx/dispatch/basic.py	80.64%	6 Missing ⚠️
pytensor/compile/mode.py	50.00%	1 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (80.23%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1365      +/-   ##
==========================================
+ Coverage   82.02%   82.04%   +0.01%     
==========================================
  Files         203      208       +5     
  Lines       48845    48949     +104     
  Branches     8691     8701      +10     
==========================================
+ Hits        40067    40162      +95     
- Misses       6627     6632       +5     
- Partials     2151     2155       +4

Files with missing lines	Coverage Δ
pytensor/link/mlx/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/mlx/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/mlx/dispatch/math.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/linker.py	`100.00% <ø> (ø)`
pytensor/compile/mode.py	`84.09% <50.00%> (-0.64%)`	⬇️
pytensor/link/mlx/dispatch/basic.py	`80.64% <80.64%> (ø)`
pytensor/link/mlx/linker.py	`77.50% <77.50%> (ø)`

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2025-04-12T12:39:06Z

I suggest basing yourself on the numba linker, torch has a lot of hacks we hopefully don't need here

williambdean · 2025-04-12T20:29:04Z

Thanks for the pointer. I simplified the one method. Do you think that gen_functors can be removed as well? The only commonality with pytorch then is that no input can be numpy array.

ricardoV94 · 2025-04-13T13:19:54Z

Yeah you shouldn't need that you just need a call to tipify on the runtime inputs as well

pytensor/compile/mode.py

pytensor/link/mlx/linker.py

ricardoV94 · 2025-04-24T09:46:39Z

pytensor/link/mlx/dispatch/basic.py

+
+
+@mlx_typify.register(np.ndarray)
+@mlx_typify.register(mx.array)


mxarray should be registered in mlx_typify_no_conversion_needed

ricardoV94 · 2025-04-24T09:47:06Z

pytensor/link/mlx/dispatch/basic.py

+@mlx_funcify.register(Assert)
+@mlx_funcify.register(CheckAndRaise)
+def mlx_funcify_CheckAndRaise(op, **kwargs):
+    warnings.warn(
+        f"""Skipping `CheckAndRaise` Op (assertion: {op.msg}) as MLX tracing would remove it.""",
+        stacklevel=2,
+    )
+
+    def assert_fn(x, *inputs):
+        return x
+
+    return assert_fn


Is this true, or just copy/pasta from JAX?

I need to check more here!

ricardoV94 · 2025-04-24T09:47:35Z

pytensor/link/mlx/dispatch/blockwise.py

+from pytensor.tensor.signal.conv import Conv1d
+
+
+def blockwise_conv1d(op, node, **kwargs):


Not needed anymore since they fixed upstream right?

I think we needed still. Where do you see its fixed? We are using this blockwise conv1d.

Sure but blockwise will call vmap on the core op, so we only need to dispatch core Conv1D to MLX Conv1D, then the blockwise variant will work automatically

ricardoV94 · 2025-04-24T09:48:08Z

pytensor/link/mlx/dispatch/core.py

+# ------------------------------------------------------------------
+# Join
+# ------------------------------------------------------------------


Don't introduce comments like this, the code is readable and won't get stale if we move things around

ricardoV94 · 2025-04-24T09:48:25Z

pytensor/link/mlx/dispatch/core.py

+# ------------------------------------------------------------------
+# Join
+# ------------------------------------------------------------------
+@mlx_funcify.register(Join)  # MLX


The multiple #MLX comments are useless

ricardoV94 · 2025-04-24T09:50:45Z

pytensor/link/mlx/dispatch/elemwise.py

+        # Convert scalar to array if needed
+        if isinstance(x, int | float) or (
+            isinstance(x, np.number) and not isinstance(x, np.ndarray)
+        ):
+            x = mx.array(x)


Should not be needed

ricardoV94 · 2025-04-24T09:51:27Z

pytensor/link/mlx/dispatch/elemwise.py

+@mlx_funcify.register(CAReduce)
+def mlx_funcify_CAReduce(op, **kwargs):
+    if isinstance(op.scalar_op, Add):
+
+        def sum(x):
+            return mx.sum(x, axis=op.axis)
+
+        return sum
+    elif isinstance(op.scalar_op, Mul):
+
+        def prod(x):
+            return mx.prod(x, axis=op.axis)
+
+        return prod
+    elif isinstance(op.scalar_op, AND):
+
+        def all(x):
+            return x.all(axis=op.axis)
+
+        return all
+    elif isinstance(op.scalar_op, OR):
+
+        def any(x):
+            return mx.any(x, axis=op.axis)
+
+        return any
+    else:
+        raise NotImplementedError(f"MLX does not support Elemwise {op.scalar_op}")


This should do a secon-lever dispatch on the core_op. Something like MLX_funcify_CAREduce that is called on op.scalar_op

ricardoV94 · 2025-04-24T09:53:14Z

pytensor/link/mlx/dispatch/math.py

+
+
+@mlx_funcify.register(Elemwise)
+def mlx_funcify_Elemwise(op, **kwargs):


Like CAReduce it should have a second level dispatch. Also we need to enforce the runtime_broadcastable checks (same in Alloc). And we shoud have a default implementation for that second level dispatch that tries to use getattr(MLX, "func_name") similar to how JAX does it already.

ricardoV94 · 2025-04-24T09:53:55Z

pytensor/link/mlx/dispatch/subtensor.py

+            if not op.inplace:
+                x = deepcopy(x)
+            x[indices] = y
+            return x


Need tests for all these including inplace variants

ricardoV94 · 2025-04-24T09:54:22Z

pytensor/link/mlx/linker.py

+        for n in self.fgraph.inputs:
+            sinput = storage_map[n]
+            # Handle random number generators specially
+            if isinstance(sinput[0], RandomState | Generator):
+                new_value = mlx_typify(
+                    sinput[0], dtype=getattr(sinput[0], "dtype", None)
+                )
+                sinput[0] = new_value
+            thunk_inputs.append(sinput)


Since we don't have Random stuff yet we shouldn't include the code

Okay added some!

ricardoV94 · 2025-06-03T08:27:35Z

pytensor/link/mlx/dispatch/blockwise.py

+    # 1) If it's a Conv1d Blockwise, use the custom implementation
+    if isinstance(op.core_op, Conv1d):
+        return blockwise_conv1d(op, node, **kwargs)


Here, we don't need this special casing anymore

Are you sure? Without it the test fails.

ricardoV94 · 2025-06-03T08:27:50Z

pytensor/link/mlx/dispatch/core.py

+def mlx_funcify_Join(op, **kwargs):
+    def join(axis, *tensors):
+        view = op.view
+        if (view != -1) and all(


view was removed from PyTensor, ignore

ricardoV94 · 2025-06-03T08:30:18Z

pytensor/link/mlx/dispatch/math.py

+@mlx_typify.register(int)
+@mlx_typify.register(float)
+def mlx_typify_python_scalar(data, **kwargs):
+    return mx.array(data)


Put this close to the other mlx_typipy

ricardoV94 · 2025-06-03T08:31:29Z

pytensor/link/mlx/dispatch/math.py

+    def elemwise(*inputs):
+        # Enforce runtime broadcast checks (same as JAX and PyTorch implementations)
+        if node is not None:
+            # Convert inputs to MLX arrays for broadcast checking


node is never none, conversion to mx.array() shouldn't be needed either

ricardoV94 · 2025-06-03T08:32:59Z

tests/link/mlx/test_shape.py

+from tests.link.mlx.test_basic import compare_mlx_and_py
+
+
+@pytest.mark.xfail(reason="Shape Op is not supported yet")


Seems like it is?

ricardoV94 · 2025-06-03T08:34:09Z

pytensor/link/mlx/dispatch/elemwise.py

+    return min_reduce
+
+
+@mlx_funcify.register(CAReduce)


Put this before the specific scalar_op dispatches?

ricardoV94 · 2025-06-03T08:34:35Z

pytensor/link/mlx/dispatch/elemwise.py

+    return softmax_grad
+
+
+@mlx_funcify.register(Softplus)


Delete this? You have one in elemwise already

williambdean · 2025-06-03T11:51:57Z

Were we going to split this PR up into core functionality and op implementations? What are the next steps here?

ricardoV94 · 2025-06-03T11:59:19Z

This PR seems okay state that we can merge as is when the comments are addressed

cetagostini · 2025-06-08T20:10:23Z

@ricardoV94 I think all it's applied.

ricardoV94 · 2025-06-09T07:01:34Z

pytensor/link/mlx/dispatch/core.py

@@ -177,7 +177,7 @@ def tensor_from_scalar(x):
 @mlx_funcify.register(ScalarFromTensor)
 def mlx_funcify_ScalarFromTensor(op, **kwargs):
    def scalar_from_tensor(x):
-        return x.reshape(-1)[0]
+        return mx.array(x).reshape(-1)[0]


This seems convoluted. Does MLX have something like x.item()?

I need to check this again. One minute.

It has but we have issues with the pytensor-mlx way to compile, and this make it generic enough.

Optimize to have both options, hope you like that more!

ricardoV94 · 2025-06-09T07:02:51Z

tests/link/mlx/test_basic.py

+    result = scalar_from_tensor_func(mlx_array)
+    assert result == 42
+
+    # Test with Python int (this used to fail)


Why are python ints being passed? That suggests a bug elsewhere

Indeed, not bug but didn't replace a few things, done!

ricardoV94 · 2025-06-09T07:03:59Z

tests/link/mlx/test_basic.py

+    scalar_result = pt.scalar_from_tensor(x)
+
+    # Create function and test
+    f = pytensor.function([], scalar_result, mode="MLX")


This won't test MLX at all because it's a constant function and you are running full optimization, so it will just be constant folded. Instead make x a symbolic variable like x = pytensor.scalar.int64(x) that is an input to the function.

Totally right, miss that 🙌🏻

ricardoV94 · 2025-06-09T07:04:23Z

tests/link/mlx/test_elemwise.py

@@ -11,3 +12,39 @@ def test_input(op) -> None:
    x_test = mx.array([1.0, 2.0, 3.0])

    compare_mlx_and_py([x], out, [x_test])
+
+
+def test_new_elemwise_operations() -> None:


bad test name, they won't be new by the time this PR is merged ;)

ricardoV94 · 2025-06-09T07:06:11Z

tests/link/mlx/test_shape.py

 def test_mlx_Reshape_constant():
    a = vector("a")
    x = reshape(a, (2, 2))
    compare_mlx_and_py([a], [x], [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)])


-@pytest.mark.xfail(reason="Reshape Op is not supported yet")
+def test_mlx_Reshape_various_shapes():


Isn't this similar to test_mlx_Reshape_concrete below? Combine with that?

I mean, doesn't seem appropriate in this case, as they're testing different scenarios of the reshape operation.

test_mlx_Reshape_various_shapes focuses on testing different dimensional transformations with static/constant shapes.

test_mlx_Reshape_concrete_shape focuses on testing computed/dynamic shapes where the shape is derived from the input tensor's properties.

Maybe they can be rename? But I feel two different things!

cetagostini · 2025-06-09T09:14:08Z

Current implementation allow to sample a simple pymc-marketing model, both gpu and cpu with MLX backend. Nevertheless complex model got issues still.

review-notebook-app · 2025-06-09T16:44:20Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

cetagostini · 2025-06-09T16:45:57Z

Adding small implementation to test MLX speed against jax.

williambdean added 2 commits April 11, 2025 11:38

mlx poc

d25f214

add test for dot

edacc0e

williambdean marked this pull request as draft April 11, 2025 15:43

williambdean commented Apr 11, 2025

View reviewed changes

pytensor/link/pytorch/linker.py Outdated Show resolved Hide resolved

williambdean added 10 commits April 11, 2025 12:26

restore pytorch

052fdc2

wrap in mx.array

a9ecad0

modify the pytorch jit

e690bff

move file

ad29c17

dont wrap

ba29b37

attempt to fix github action

8716870

change the rtol

9bf7edf

add init file

96ba116

skip if not installed

e116fa1

remove torch related code / comments

5d5f754

williambdean marked this pull request as ready for review April 11, 2025 17:54

williambdean added 2 commits April 12, 2025 16:25

simplify the fgraph_convert

b8cee3f

assert type

d057453

ricardoV94 reviewed Apr 13, 2025

View reviewed changes

pytensor/compile/mode.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Apr 13, 2025

View reviewed changes

pytensor/link/mlx/linker.py Outdated Show resolved Hide resolved

williambdean and others added 7 commits April 18, 2025 13:51

simplify the internal

ae202e6

remove the language

f1941fe

Adding operations in pytensor

7c8eae7

add extension

67a74fb

make compare function

fb5eb52

rename function

516b595

correct the function name

67bb8da

ricardoV94 requested changes Apr 24, 2025

View reviewed changes

ricardoV94 added enhancement New feature or request backend compatibility labels Apr 24, 2025

cetagostini mentioned this pull request Apr 28, 2025

Computational overhead when transforming to a numpy array ml-explore/mlx#2132

Closed

cetagostini added 2 commits June 2, 2025 20:45

Requested changes by Ricardo

6a2b774

Pre commit changes

602f0ed

ricardoV94 reviewed Jun 3, 2025

View reviewed changes

cetagostini added 2 commits June 8, 2025 23:08

More changes from Ricardo

8a2aea9

Pre Commit RUN

845561c

Adding more operations for complex model

6ab7428

ricardoV94 reviewed Jun 9, 2025

View reviewed changes

Working with simple model

b6292f1

cetagostini added 6 commits June 9, 2025 12:19

Change bad name

f2d9d1b

Correcting test by Ricardo

5c759b9

Changing synth test

97b2e31

Optimizing reshape

dd83e0f

Comment

662b4f2

Small changes and adding small benchmark

bcf7f8d



		@mlx_typify.register(np.ndarray)
		@mlx_typify.register(mx.array)

		from pytensor.tensor.signal.conv import Conv1d


		def blockwise_conv1d(op, node, **kwargs):



		@mlx_funcify.register(Elemwise)
		def mlx_funcify_Elemwise(op, **kwargs):

		from tests.link.mlx.test_basic import compare_mlx_and_py


		@pytest.mark.xfail(reason="Shape Op is not supported yet")

MLX backend POC #1365

Are you sure you want to change the base?

MLX backend POC #1365

Conversation

williambdean commented Apr 11, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

Uh oh!

codecov bot commented Apr 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 commented Apr 12, 2025

Uh oh!

williambdean commented Apr 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Apr 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

williambdean commented Apr 11, 2025 •

edited by github-actions bot

Loading

codecov bot commented Apr 11, 2025 •

edited

Loading

williambdean commented Apr 12, 2025 •

edited

Loading

ricardoV94 commented Apr 13, 2025 •

edited

Loading

ricardoV94 commented Jun 3, 2025 •

edited

Loading