Replace dot of 1 and x -> x

pytensor/tensor/rewriting/math.py

@@ -190,6 +190,43 @@ def local_0_dot_x(fgraph, node):
  return [constant_zero]
 
 
+@register_canonicalize
+@register_stabilize
+@node_rewriter([Dot])
+def local_1_dot_x(fgraph, node):
+ if not isinstance(node.op, Dot):
+ return False
+
+ x = node.inputs[0]
+ y = node.inputs[1]
+ replace = False
+ try:
+ if get_underlying_scalar_constant_value(x, only_process_constants=True) == 1:
+ replace = True
+ var = y
+ except NotScalarConstantError:
+ pass
+
+ try:
+ if get_underlying_scalar_constant_value(y, only_process_constants=True) == 1:
+ replace = True
+ var = x
+ except NotScalarConstantError:
+ pass
+
+ if replace:
+ new_out = var
+ old_out = node.outputs[0]
+
+ if new_out.dtype != old_out.dtype:
+ new_out = cast(new_out, old_out.dtype)
+ if not old_out.type.is_super(new_out.type):
+ new_out = new_out.reshape(old_out.shape)
+ # new_out = alloc_like(new_out, old_out, fgraph)
+
+ return [new_out]
+
+
 @register_canonicalize
 @node_rewriter([DimShuffle])
 def local_lift_transpose_through_dot(fgraph, node):

tests/tensor/rewriting/test_math.py

@@ -4475,3 +4475,22 @@ def test_local_batched_matmul_to_core_matmul():
  x_test = rng.normal(size=(5, 3, 2))
  y_test = rng.normal(size=(5, 2, 2))
  np.testing.assert_allclose(fn(x_test, y_test), x_test @ y_test)
+
+
+@pytest.mark.parametrize(
+ "x",
+ (
+ pt.col("x"),
+ fmatrix("x"),
+ vector("x"),
+ pt.tensor("x", shape=(1, 3, 2), dtype="float64"),
+ ),
+)
+def test_mul_with_1(x):
+ f = x @ [[1.0]]
+ with pytensor.config.change_flags(optimizer_verbose=True):
+ fn = pytensor.function([x], f, mode=get_default_mode().excluding("BlasOpt"))
+
+ pytensor.dprint(fn)
+
+ assert 0