Optimizing reshape

Author: cetagostini

pytensor/link/mlx/dispatch/core.py

@@ -212,7 +212,17 @@ def tensor_from_scalar(x):
 @mlx_funcify.register(ScalarFromTensor)
 def mlx_funcify_ScalarFromTensor(op, **kwargs):
     def scalar_from_tensor(x):
-        return mx.array(x).reshape(-1)[0]
+        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
 
     return scalar_from_tensor
 

tests/link/mlx/test_basic.py

@@ -1,6 +1,7 @@
 """
 Basic tests for the MLX backend.
 """
+
 from collections.abc import Callable, Iterable
 from functools import partial