almost there baby william

Author: cetagostini

pytensor/link/mlx/dispatch/basic.py

@@ -1,3 +1,4 @@
+import warnings
 from functools import singledispatch
 from types import NoneType
 
@@ -7,6 +8,7 @@
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.fg import FunctionGraph
 from pytensor.link.utils import fgraph_to_python
+from pytensor.raise_op import Assert, CheckAndRaise
 
 
 @singledispatch
@@ -59,3 +61,17 @@ def deepcopyop(x):
         return x.copy()
 
     return deepcopyop
+
+
+@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

pytensor/link/mlx/dispatch/blockwise.py

@@ -1,66 +1,18 @@
 import mlx.core as mx
 
-from pytensor.graph import FunctionGraph
 from pytensor.link.mlx.dispatch import mlx_funcify
 from pytensor.tensor.blockwise import Blockwise
 
 
 @mlx_funcify.register(Blockwise)
 def funcify_Blockwise(op: Blockwise, node, *args, **kwargs):
-    # Create a function graph for the core operation
     core_node = op._create_dummy_core_node(node.inputs)
-    core_fgraph = FunctionGraph(inputs=core_node.inputs, outputs=core_node.outputs)
+    core_f = mlx_funcify(op.core_op, core_node)
+    blockwise_f = core_f
+    for i in range(op.batch_ndim(node)):
+        blockwise_f = mx.vmap(blockwise_f)
 
-    # Convert the core function graph to an MLX function
-    tuple_core_fn = mlx_funcify(core_fgraph, **kwargs)
+    def blockwise_fun(*inputs):
+        return blockwise_f(*inputs)
 
-    # If there's only one output, unwrap it from the tuple
-    if len(node.outputs) == 1:
-
-        def core_fn(*inputs):
-            return tuple_core_fn(*inputs)[0]
-    else:
-        core_fn = tuple_core_fn
-
-    # Apply vmap for each batch dimension
-    batch_ndims = op.batch_ndim(node)
-    vmap_fn = core_fn
-    for _ in range(batch_ndims):
-        vmap_fn = mx.vmap(vmap_fn)
-
-    def blockwise_fn(*inputs):
-        # Check for runtime broadcasting compatibility
-        op._check_runtime_broadcast(node, inputs)
-
-        # Handle broadcasting for batched dimensions
-        if batch_ndims > 0:
-            # Get batch shapes for broadcasting
-            batch_shapes = [inp.shape[:batch_ndims] for inp in inputs]
-
-            # Calculate the broadcasted batch shape
-            from functools import reduce
-
-            def broadcast_shapes(shape1, shape2):
-                return tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2, strict=True))
-
-            if batch_shapes:
-                broadcasted_shape = reduce(broadcast_shapes, batch_shapes)
-
-                # Broadcast inputs to the common batch shape
-                broadcasted_inputs = []
-                for inp in inputs:
-                    if inp.shape[:batch_ndims] != broadcasted_shape:
-                        # Create the full target shape
-                        target_shape = broadcasted_shape + inp.shape[batch_ndims:]
-                        # Broadcast the input
-                        broadcasted_inputs.append(mx.broadcast_to(inp, target_shape))
-                    else:
-                        broadcasted_inputs.append(inp)
-
-                # Apply the vectorized function to the broadcasted inputs
-                return vmap_fn(*broadcasted_inputs)
-
-        # No broadcasting needed
-        return vmap_fn(*inputs)
-
-    return blockwise_fn
+    return blockwise_fun

pytensor/link/mlx/dispatch/math.py

@@ -1,36 +1,40 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch import mlx_funcify
+from pytensor.scalar import Softplus
 from pytensor.scalar.basic import (
+    AND,
     EQ,
     GE,
     GT,
     LE,
     LT,
     NEQ,
+    OR,
     Abs,
     Add,
+    Cast,
     Cos,
     Exp,
     Log,
     Mul,
+    Neg,
     Pow,
+    ScalarMaximum,
+    ScalarMinimum,
+    Sign,
     Sin,
     Sqr,
     Sqrt,
     Sub,
     Switch,
     TrueDiv,
-    Neg,
-    AND,
-    OR,
-    ScalarMaximum,
-    ScalarMinimum,
+    Log1p
 )
 from pytensor.scalar.math import Sigmoid
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.math import Dot
-from pytensor.scalar import Softplus
+
 
 @mlx_funcify.register(Dot)
 def mlx_funcify_Dot(op, **kwargs):
@@ -169,6 +173,7 @@ def abs(x):
 
         return abs
     elif isinstance(op.scalar_op, Softplus):
+
         def softplus(x):
             return mx.where(
                 x < -37.0,
@@ -194,7 +199,7 @@ def neg(x):
     elif isinstance(op.scalar_op, AND):
 
         def all(x):
-            return mx.all(x, axis=op.axis)
+            return mx.all(x)
 
         return all
     elif isinstance(op.scalar_op, OR):
@@ -215,5 +220,23 @@ def min(x):
             return mx.min(x, axis=op.axis)
 
         return min
+    elif isinstance(op.scalar_op, Cast):
+
+        def cast(x):
+            return mx.cast(x, op.dtype)
+
+        return cast
+    elif isinstance(op.scalar_op, Sign):
+
+        def sign(x):
+            return mx.sign(x)
+
+        return sign
+    elif isinstance(op.scalar_op, Log1p):
+
+        def log1p(x):
+            return mx.log1p(x)
+
+        return log1p
     else:
         raise NotImplementedError(f"MLX does not support {op.scalar_op}")

pytensor/link/mlx/dispatch/shape.py

@@ -18,7 +18,7 @@ def specifyshape(x, *shape):
 
 @mlx_funcify.register(Shape_i)
 def mlx_funcify_Shape_i(op, node, **kwargs):
-    def shape_i(x, i):
+    def shape_i(x):
         return x.shape[op.i]
 
     return shape_i

pytensor/link/mlx/dispatch/subtensor.py

@@ -11,40 +11,32 @@
 from pytensor.tensor.type_other import MakeSlice
 
 
-BOOLEAN_MASK_ERROR = """MLX does not support resizing arrays with boolean
-masks. In some cases, however, it is possible to re-express your model
-in a form that MLX can compile:
-
->>> import pytensor.tensor as pt
->>> x_pt = pt.vector('x')
->>> y_pt = x_pt[x_pt > 0].sum()
-
-can be re-expressed as:
+@mlx_funcify.register(Subtensor)
+def mlx_funcify_Subtensor(op, node, **kwargs):
+    idx_list = getattr(op, "idx_list", None)
 
->>> import pytensor.tensor as pt
->>> x_pt = pt.vector('x')
->>> y_pt = pt.where(x_pt > 0, x_pt, 0).sum()
-"""
+    def subtensor(x, *ilists):
+        indices = indices_from_subtensor([int(element) for element in ilists], idx_list)
+        if len(indices) == 1:
+            indices = indices[0]
 
-DYNAMIC_SLICE_LENGTH_ERROR = """MLX does not support slicing arrays with a dynamic
-slice length.
-"""
+        return x.__getitem__(indices)
 
+    return subtensor
 
-@mlx_funcify.register(Subtensor)
 @mlx_funcify.register(AdvancedSubtensor)
 @mlx_funcify.register(AdvancedSubtensor1)
-def mlx_funcify_Subtensor(op, node, **kwargs):
+def mlx_funcify_AdvancedSubtensor(op, node, **kwargs):
     idx_list = getattr(op, "idx_list", None)
 
-    def subtensor(x, *ilists):
+    def advanced_subtensor(x, *ilists):
         indices = indices_from_subtensor(ilists, idx_list)
         if len(indices) == 1:
             indices = indices[0]
 
         return x.__getitem__(indices)
 
-    return subtensor
+    return advanced_subtensor
 
 
 @mlx_funcify.register(IncSubtensor)