pre-commit

Author: cetagostini

pytensor/link/mlx/dispatch/blockwise.py

@@ -7,11 +7,11 @@
 @mlx_funcify.register(Blockwise)
 def funcify_Blockwise(op: Blockwise, node, **kwargs):
     """Convert a Blockwise operation to an MLX function.
-    
+
     This handles vectorized operations by using mx.vmap when there are batch
     dimensions, or returning the core function directly when there are no
     batch dimensions to vectorize over.
-    
+
     Parameters
     ----------
     op : Blockwise
@@ -20,7 +20,7 @@ def funcify_Blockwise(op: Blockwise, node, **kwargs):
         The node containing the operation and its inputs.
     **kwargs
         Additional keyword arguments.
-        
+
     Returns
     -------
     callable

tests/link/mlx/test_blockwise.py

@@ -29,41 +29,41 @@ def test_blockwise_conv1d():
 
 def test_blockwise_no_batch_dimensions():
     """Test that Blockwise returns the core function when there are no batch dimensions.
-    
+
     This verifies the fix for the vmap dispatcher issue where mx.vmap should not
     be called when there are no batch dimensions to vectorize over.
     """
     rng = np.random.default_rng(42)
-    
+
     # Create a blockwise matmul with no batch dimensions (core operation only)
     x = pt.matrix("x")
     y = pt.matrix("y")
-    
+
     blockwise_matmul = Blockwise(Dot(), signature="(i,j),(j,k)->(i,k)")
     z = blockwise_matmul(x, y)
-    
+
     x_test = rng.normal(size=(2, 3))
     y_test = rng.normal(size=(3, 4))
-    
+
     compare_mlx_and_py([x, y], [z], [x_test, y_test], must_be_device_array=True)
 
 
 def test_blockwise_all_broadcastable_batch_dims():
     """Test that Blockwise returns the core function when all batch dims are broadcastable.
-    
+
     When all batch dimensions are size-1 (broadcastable), vmap should not be called
     since there's no actual vectorization needed.
     """
     rng = np.random.default_rng(43)
-    
+
     # Create inputs with size-1 batch dimensions
     x = tensor("x", shape=(1, 2, 3))
     y = tensor("y", shape=(1, 3, 4))
-    
+
     blockwise_matmul = Blockwise(Dot(), signature="(i,j),(j,k)->(i,k)")
     z = blockwise_matmul(x, y)
-    
+
     x_test = rng.normal(size=(1, 2, 3))
     y_test = rng.normal(size=(1, 3, 4))
-    
+
     compare_mlx_and_py([x, y], [z], [x_test, y_test], must_be_device_array=True)