Rewrite scalar solve to division

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -1020,3 +1020,29 @@ def slogdet_specialization(fgraph, node):
             k: slogdet_specialization_map[v] for k, v in dummy_replacements.items()
         }
         return replacements
+
+
+@register_stabilize
+@register_canonicalize
+@node_rewriter([Blockwise])
+def scalar_solve_to_divison(fgraph, node):
+    """
+    Replace solve(a, b) with b / a if a is a (1, 1) matrix
+    """
+
+    core_op = node.op.core_op
+    if not isinstance(core_op, Solve):
+        return None
+
+    a, b = node.inputs
+    old_out = node.outputs[0]
+    if not all(a.broadcastable[-2:]):
+        return None
+
+    new_out = b / a
+    if core_op.b_ndim == 1:
+        new_out = new_out.squeeze(-1)
+
+    copy_stack_trace(old_out, new_out)
+
+    return [new_out]

tests/tensor/rewriting/test_linalg.py

@@ -993,3 +993,25 @@ def test_slogdet_specialization():
     f = function([x], [exp_det_x, sign_det_x], mode="FAST_RUN")
     nodes = f.maker.fgraph.apply_nodes
     assert not any(isinstance(node.op, SLogDet) for node in nodes)
+
+
+def test_scalar_solve_to_division_rewrite():
+    rng = np.random.default_rng(sum(map(ord, "scalar_solve_to_division_rewrite")))
+
+    a = pt.dmatrix("a", shape=(1, 1))
+    b = pt.dvector("b")
+
+    c = pt.linalg.solve(a, b, b_ndim=1)
+
+    f = function([a, b], c, mode="FAST_RUN")
+    nodes = f.maker.fgraph.apply_nodes
+
+    assert not any(isinstance(node.op, Solve) for node in nodes)
+
+    a_val = rng.normal(size=(1, 1)).astype(pytensor.config.floatX)
+    b_val = rng.normal(size=(1,)).astype(pytensor.config.floatX)
+
+    c_val = np.linalg.solve(a_val, b_val)
+    np.testing.assert_allclose(
+        f(a_val, b_val), c_val, rtol=1e-7 if config.floatX == "float64" else 1e-5
+    )