Refactor sum_prod_mul rewrite test and add failing case

Rewrite from prod of mul was not correct when only some axes were reduced by prod

Author: ricardoV94

tests/tensor/rewriting/test_math.py

@@ -2512,6 +2512,7 @@ def test_local_sum_prod_mul_by_scalar(self):
         # 4-the inputs to the mul contain two scalars and no non-scalar
         # 5-the inputs to the mul contain two scalars and one non-scalar
         # 6-the inputs to the mul contain two scalars and two non-scalars
+        # 7-the reduction happens across only the first of two axes
 
         vect = dvector()
         mat = dmatrix()
@@ -2524,10 +2525,15 @@ def test_local_sum_prod_mul_by_scalar(self):
         s2_val = np.random.random()
 
         def test_reduction_rewrite(
-            inputs, inputs_val, reduction_op, expected_output, nb_expected_sum_nodes
+            inputs,
+            inputs_val,
+            reduction_op,
+            expected_output,
+            nb_expected_sum_nodes,
+            axis=None,
         ):
             mul_out = mul(*inputs)
-            f = function(inputs, reduction_op()(mul_out), mode=self.mode)
+            f = function(inputs, reduction_op(axis=axis)(mul_out), mode=self.mode)
             out = f(*inputs_val)
             utt.assert_allclose(out, expected_output)
 
@@ -2581,6 +2587,16 @@ def test_reduction_rewrite(
             1,
         )
 
+        # Case 7
+        test_reduction_rewrite(
+            [mat, scalar1, scalar2],
+            [m_val, s1_val, s2_val],
+            Sum,
+            (s1_val * s2_val * m_val).sum(0),
+            1,
+            axis=(0,),
+        )
+
         # Test prod
 
         # Case 1
@@ -2627,6 +2643,16 @@ def test_reduction_rewrite(
             2,
         )
 
+        # Case 7
+        test_reduction_rewrite(
+            [mat, scalar1, scalar2],
+            [m_val, s1_val, s2_val],
+            Prod,
+            (s1_val * s2_val * m_val).prod(0),
+            1,
+            axis=(0,),
+        )
+
     def test_local_sum_prod_all_to_none(self):
         a = tensor3()
         input = np.arange(3 * 4 * 5, dtype=config.floatX).reshape(3, 4, 5)