Changes with Ricardo

Author: cetagostini

doc/_drafts/benchmark_mlx_v_jax_corrected.ipynb

@@ -2,7 +2,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": 15,
+      "execution_count": 1,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -21,7 +21,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 23,
+      "execution_count": 2,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -296,7 +296,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 24,
+      "execution_count": null,
       "metadata": {},
       "outputs": [
         {
@@ -306,16 +306,14 @@
             "Verifying computational correctness...\n",
             "Max difference between JAX and MLX results: 0.00e+00\n",
             "✅ Results match within tolerance\n",
-            "Running benchmarks with N=1000 repetitions per test...\n",
-            "Testing 128x128 matrices...\n",
-            "Testing 256x256 matrices...\n",
-            "Testing 512x512 matrices...\n",
-            "Testing 1024x1024 matrices...\n"
+            "Running benchmarks with N=20 repetitions per test...\n",
+            "Testing 128x128 matrices...\n"
           ]
         }
       ],
       "source": [
-        "_, results = main()"
+        "iteration=20\n",
+        "_, results = main(N=iteration)"
       ]
     },
     {
@@ -346,7 +344,7 @@
         }
       ],
       "source": [
-        "print(\"\\nBenchmark Results over 1000 repetitions:\")\n",
+        "print(f\"\\nBenchmark Results over {iteration} repetitions:\")\n",
         "print(results.to_string(index=False))"
       ]
     },

pytensor/link/mlx/dispatch/core.py

@@ -212,17 +212,8 @@ def tensor_from_scalar(x):
 @mlx_funcify.register(ScalarFromTensor)
 def mlx_funcify_ScalarFromTensor(op, **kwargs):
     def scalar_from_tensor(x):
-        arr = mx.array(x)
-        try:
-            # Try .item() first (cleaner and faster when possible)
-            return arr.item()
-        except ValueError as e:
-            if "eval" in str(e):
-                # Fall back to reshape approach for compiled contexts
-                return arr.reshape(-1)[0]
-            else:
-                # Re-raise if it's a different error
-                raise
+        "We can't not return a scalar in MLX without trigger evaluation"
+        return x
 
     return scalar_from_tensor
 

pytensor/link/mlx/dispatch/shape.py

@@ -7,7 +7,7 @@
 @mlx_funcify.register(Shape)
 def mlx_funcify_Shape(op, **kwargs):
     def shape(x):
-        return x.shape
+        return mx.array(x.shape, dtype=mx.int64)
 
     return shape
 

tests/link/mlx/test_basic.py

@@ -1,6 +1,7 @@
 """
 Basic tests for the MLX backend.
 """
+import pytest
 
 from collections.abc import Callable, Iterable
 from functools import partial
@@ -224,12 +225,7 @@ def test_alloc_pytensor_integration():
     x = pt.scalar("x", dtype="float32")
     result = pt.alloc(x, 3, 4)
 
-    # Use MLX mode
-    from pytensor.compile import mode
-
-    mlx_mode = mode.get_mode("MLX")
-
-    f = pytensor.function([x], result, mode=mlx_mode)
+    f = pytensor.function([x], result, mode="MLX")
     output = f(5.0)
 
     assert output.shape == (3, 4)
@@ -238,7 +234,6 @@ def test_alloc_pytensor_integration():
 
 def test_alloc_compilation_limitation():
     """Test that Alloc operations with dynamic shapes provide helpful error in compiled contexts."""
-    import pytest
 
     # Create variables
     x = pt.scalar("x", dtype="float32")