Use Blockwise by default

Author: ricardoV94

pytensor/tensor/blockwise.py

@@ -171,13 +171,6 @@ def make_node(self, *inputs):
             inp.type.ndim - len(sig) for inp, sig in zip(inputs, self.inputs_sig)
         )
 
-        # Don't pollute the graph with useless BlockWise
-        # TODO: Do we want to do this? Or leave it as a Blockwise and later have a rewrite that removes useless casse
-        #  A reason to not eagerly avoid Blockwise is that we could make all rewrites track the Blockwise version,
-        #  instead of having to track both or only the more restricted core case.
-        if not batch_ndims:
-            return self.core_op.make_node(*inputs)
-
         batched_inputs = []
         batch_shapes = []
         for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig)):

pytensor/tensor/nlinalg.py

@@ -10,6 +10,7 @@
 from pytensor.tensor import basic as at
 from pytensor.tensor import math as tm
 from pytensor.tensor.basic import as_tensor_variable, extract_diag
+from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.type import dvector, lscalar, matrix, scalar, vector
 
 
@@ -76,7 +77,7 @@ def pinv(x, hermitian=False):
     solve op.
 
     """
-    return MatrixPinv(hermitian=hermitian)(x)
+    return Blockwise(MatrixPinv(hermitian=hermitian))(x)
 
 
 class MatrixInverse(Op):
@@ -153,7 +154,7 @@ def infer_shape(self, fgraph, node, shapes):
         return shapes
 
 
-inv = matrix_inverse = MatrixInverse()
+inv = matrix_inverse = Blockwise(MatrixInverse())
 
 
 def matrix_dot(*args):
@@ -214,7 +215,7 @@ def __str__(self):
         return "Det"
 
 
-det = Det()
+det = Blockwise(Det())
 
 
 class SLogDet(Op):
@@ -248,7 +249,7 @@ def __str__(self):
         return "SLogDet"
 
 
-slogdet = SLogDet()
+slogdet = Blockwise(SLogDet())
 
 
 class Eig(Op):
@@ -278,7 +279,7 @@ def infer_shape(self, fgraph, node, shapes):
         return [(n,), (n, n)]
 
 
-eig = Eig()
+eig = Blockwise(Eig())
 
 
 class Eigh(Eig):

pytensor/tensor/rewriting/__init__.py

@@ -2,6 +2,7 @@
 import pytensor.tensor.rewriting.blas
 import pytensor.tensor.rewriting.blas_c
 import pytensor.tensor.rewriting.blas_scipy
+import pytensor.tensor.rewriting.blockwise
 import pytensor.tensor.rewriting.elemwise
 import pytensor.tensor.rewriting.extra_ops
 

pytensor/tensor/rewriting/blockwise.py

@@ -0,0 +1,39 @@
+from pytensor.compile.mode import optdb
+from pytensor.graph import node_rewriter
+from pytensor.graph.rewriting.basic import copy_stack_trace, out2in
+from pytensor.tensor.blockwise import Blockwise, vectorize_node
+from pytensor.tensor.rewriting.basic import register_useless
+
+
+@register_useless("fast_compile")
+@node_rewriter([Blockwise])
+def local_useless_blockwise(fgraph, node):
+    # If there is a dispatch implementation that does not require Blockwise, use that instead.
+    # This means a user created a Blockwise manually when there was no need.
+    op = node.op
+    inputs = node.inputs
+    dummy_core_node = op._create_dummy_core_node(node.inputs)
+    vect_node = vectorize_node(dummy_core_node, *inputs)
+    if not isinstance(vect_node.op, Blockwise):
+        return copy_stack_trace(node.outputs, vect_node.outputs)
+
+
+@node_rewriter([Blockwise])
+def local_useless_unbatched_blockwise(fgraph, node):
+    """Remove Blockwise that don't have any batched dims."""
+    op = node.op
+    inputs = node.inputs
+
+    if max(inp.type.ndim - len(sig) for inp, sig in zip(inputs, op.inputs_sig)) == 0:
+        return copy_stack_trace(node.outputs, op.core_op.make_node(*inputs).outputs)
+
+
+# We register this rewrite late, so that other rewrites need only target Blockwise Ops
+optdb.register(
+    "local_useless_unbatched_blockwise",
+    out2in(local_useless_unbatched_blockwise, ignore_newtrees=True),
+    "fast_run",
+    "fast_compile",
+    "blockwise",
+    position=49,
+)

pytensor/tensor/utils.py

@@ -1,5 +1,3 @@
-from typing import Callable, Optional
-
 import numpy as np
 
 import pytensor
@@ -132,4 +130,4 @@ def import_func_from_string(func_string: str):  # -> Optional[Callable]:
             except AttributeError:
                 module = None
                 break
-        return module
+        return module

tests/tensor/rewriting/test_blockwise.py

@@ -0,0 +1,35 @@
+from pytensor import function
+from pytensor.scalar import log as scalar_log
+from pytensor.tensor import matrix, tensor3
+from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.nlinalg import MatrixPinv, pinv
+
+
+def test_useless_blockwise_of_elemwise():
+    x = matrix("x")
+    out = Blockwise(Elemwise(scalar_log), signature="()->()")(x)
+
+    assert isinstance(out.owner.op, Blockwise)
+    assert isinstance(out.owner.op.core_op, Elemwise)
+
+    fn = function([x], out, mode="FAST_COMPILE")
+    assert isinstance(fn.maker.fgraph.outputs[0].owner.op, Elemwise)
+
+
+def test_useless_unbatched_blockwise():
+    x = matrix("x")
+    out = pinv(x)
+
+    assert isinstance(out.owner.op, Blockwise)
+    assert isinstance(out.owner.op.core_op, MatrixPinv)
+
+    fn = function([x], out, mode="FAST_COMPILE")
+    assert isinstance(fn.maker.fgraph.outputs[0].owner.op, MatrixPinv)
+
+    # Test that it's not removed when there are batched dims
+    x = tensor3("x")
+    out = pinv(x)
+    fn = function([x], out, mode="FAST_COMPILE")
+    assert isinstance(fn.maker.fgraph.outputs[0].owner.op, Blockwise)
+    assert isinstance(fn.maker.fgraph.outputs[0].owner.op.core_op, MatrixPinv)

tests/tensor/test_blockwise.py

@@ -77,21 +77,6 @@ def test_vectorize_node():
     assert new_vect_node.inputs[0] is tns4
 
 
-def test_useless_blockwise():
-    cop = MatrixInverse()
-    bop = Blockwise(cop, signature=("(m, m) -> (m, m)"))
-
-    inp = tensor(shape=(None, None, None))
-    out = bop(inp)
-    assert out.owner.op is bop
-    assert out.owner.inputs[0] is inp
-
-    inp = tensor(shape=(None, None))
-    out = bop(inp)
-    assert out.owner.op is cop
-    assert out.owner.inputs[0] is inp
-
-
 class TestOp(Op):
     def make_node(self, *inputs):
         return Apply(self, inputs, [i.type() for i in inputs])