MLX backend POC

[williambdean]

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

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

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

[codecov]

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]

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

[williambdean]

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]

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

[williambdean]

Still need to get this to run:

import pytensor

pytensor.config.mode = "MLX"

[cetagostini]

Hey big thanks to @jessegrabowski and @ricardoV94 to help with this PR!

I feel the PR is huge enough. Should we make a first merge and start to iterate on next versions? Cleaning and making all more consistent with other backends.

Thanks to @williambdean to open the PR!