Add MLX backend#1365

Merged

ricardoV94 merged 88 commits intopymc-devs:mainfrom

williambdean:mlx-poc

Oct 9, 2025

Commits on Oct 8, 2025

mlx poc

williambdean
authored and
jessegrabowski
committed
add test for dot

williambdean
authored and
jessegrabowski
committed
restore pytorch

williambdean
authored and
jessegrabowski
committed
wrap in mx.array

williambdean
authored and
jessegrabowski
committed
modify the pytorch jit

williambdean
authored and
jessegrabowski
committed
move file

williambdean
authored and
jessegrabowski
committed
dont wrap

williambdean
authored and
jessegrabowski
committed
attempt to fix github action

williambdean
authored and
jessegrabowski
committed
change the rtol

williambdean
authored and
jessegrabowski
committed
add init file

williambdean
authored and
jessegrabowski
committed
skip if not installed

williambdean
authored and
jessegrabowski
committed
remove torch related code / comments

williambdean
authored and
jessegrabowski
committed
simplify the fgraph_convert

williambdean
authored and
jessegrabowski
committed
assert type

williambdean
authored and
jessegrabowski
committed
simplify the internal

williambdean
authored and
jessegrabowski
committed
remove the language

williambdean
authored and
jessegrabowski
committed
Adding operations in pytensor

cetagostini
authored and
jessegrabowski
committed
add extension

williambdean
authored and
jessegrabowski
committed
make compare function

williambdean
authored and
jessegrabowski
committed
rename function

williambdean
authored and
jessegrabowski
committed
correct the function name

williambdean
authored and
jessegrabowski
committed
tests for elemwise

williambdean
authored and
jessegrabowski
committed
Changes

cetagostini
authored and
jessegrabowski
committed
Toma tu tomate William

cetagostini
authored and
jessegrabowski
committed
Pushing changes with the core shit.

cetagostini
authored and
jessegrabowski
committed
add more tests

williambdean
authored and
jessegrabowski
committed
additional tests

williambdean
authored and
jessegrabowski
committed
test for switch with mlx

williambdean
authored and
jessegrabowski
committed
Pushing code

cetagostini
authored and
jessegrabowski
committed
Changes

cetagostini
authored and
jessegrabowski
committed
A lot of new code

cetagostini
authored and
jessegrabowski
committed
almost there baby william

cetagostini
authored and
jessegrabowski
committed
Another push small

cetagostini
authored and
jessegrabowski
committed
fix for all

williambdean
authored and
jessegrabowski
committed
fix for carlos

williambdean
authored and
jessegrabowski
committed
just return the compiled func

williambdean
authored and
jessegrabowski
committed
A change for willy may!

cetagostini
authored and
jessegrabowski
committed
FINALLY BABY LETS PARTY! (IF YOU ARE READING THIS MAKE MORE PRs)

cetagostini
authored and
jessegrabowski
committed
THE SUPER BLOCKWISEE YA YA YA YA JUUUUU

cetagostini
authored and
jessegrabowski
committed
refactor to use getattr

williambdean
authored and
jessegrabowski
committed
bring argmax test

williambdean
authored and
jessegrabowski
committed
use deepcopy

williambdean
authored and
jessegrabowski
committed
move some tests

williambdean
authored and
jessegrabowski
committed
Guys, I'm getting sad. We need help yisus!!!!!

cetagostini
authored and
jessegrabowski
committed
WILLIAM YOU NEED TO GO ANOTHER MILE! GO ON MY MATEEEEEEE, GO PHILLIES!

cetagostini
authored and
jessegrabowski
committed
RETURN, WHAT A SHAME! Sad times are coming.

cetagostini
authored and
jessegrabowski
committed
AI COULD BE COOL? OR WE ARE JUST FUCKING AROUND?

cetagostini
authored and
jessegrabowski
committed
AI RULES BABY MY MATE

cetagostini
authored and
jessegrabowski
committed
I'm going for pizzas, it was an incredible day!

cetagostini
authored and
jessegrabowski
committed
test conv1d case

williambdean
authored and
jessegrabowski
committed
SUUUUUUUUU!!!!!! LIFE IS GOING WELL. MLX FOR MEDIA MIX MODELS BAY

committed
pre-commit

cetagostini
authored and
jessegrabowski
committed
Almost working

cetagostini
authored and
jessegrabowski
committed
Last PR sampling working

cetagostini
authored and
jessegrabowski
committed
Requested changes by Ricardo

cetagostini
authored and
jessegrabowski
committed
Pre commit changes

cetagostini
authored and
jessegrabowski
committed
More changes from Ricardo

cetagostini
authored and
jessegrabowski
committed
Pre Commit RUN

cetagostini
authored and
jessegrabowski
committed
Adding more operations for complex model

cetagostini
authored and
jessegrabowski
committed
Working with simple model

cetagostini
authored and
jessegrabowski
committed
Change bad name

cetagostini
authored and
jessegrabowski
committed
Correcting test by Ricardo

cetagostini
authored and
jessegrabowski
committed
Changing synth test

cetagostini
authored and
jessegrabowski
committed
Optimizing reshape

cetagostini
authored and
jessegrabowski
committed
Comment

cetagostini
authored and
jessegrabowski
committed
Small changes and adding small benchmark

cetagostini
authored and
jessegrabowski
committed
Changes with Ricardo

cetagostini
authored and
jessegrabowski
committed
improving benchmark

cetagostini
authored and
jessegrabowski
committed
pre commit

cetagostini
authored and
jessegrabowski
committed
benchs

cetagostini
authored and
jessegrabowski
committed
Changes on the branch

cetagostini
authored and
jessegrabowski
committed
Feedback from Ricardo

cetagostini
authored and
jessegrabowski
committed
update test based on llm recommendation

cetagostini
authored and
jessegrabowski
committed
Streamline Blockwise impl
jessegrabowski
committed
clean up imports
jessegrabowski
committed
adjust github test.yml
jessegrabowski
committed
adjust github test.yml
jessegrabowski
committed
skip mlx tests in benchmark ci
jessegrabowski
committed
Absolute imports
jessegrabowski
committed
Use `importorskip` in mlx tests
jessegrabowski
committed
address feedback
jessegrabowski
committed