Fix gemm/gemv fusion under cast_policy=numpy+floatX and make tests policy-aware

pytensor/scalar/basic.py

@@ -208,10 +208,12 @@ class autocast_float_as:
 
     Examples
     --------
+    >>> from pytensor import config
     >>> from pytensor.tensor import fvector
-    >>> with autocast_float_as("float32"):
-    ...     assert (fvector() + 1.1).dtype == "float32"  # temporary downcasting
-    >>> assert (fvector() + 1.1).dtype == "float64"  # back to default behaviour
+    >>> with config.change_flags(cast_policy="custom"):
+    ...     with autocast_float_as("float32"):
+    ...         assert (fvector() + 1.1).dtype == "float32"  # temporary downcasting
+    ...     assert (fvector() + 1.1).dtype == "float64"  # default 'custom' behaviour
 
     """
 

pytensor/tensor/rewriting/blas.py

@@ -355,7 +355,13 @@ def _gemm_from_factored_list(fgraph, lst):
         # sM can be a tuple of 2 elements or an PyTensor variable.
         if isinstance(sM, tuple):
             sm0, sm1 = sM
-            sm0 = ptb.as_tensor_variable(sm0)
+            if isinstance(sm0, int | float):
+                # Python literal scales (the +/-1.0 seeded by
+                # `_gemm_canonicalize`) adopt the matrix dtype instead of
+                # following the autocast policy
+                sm0 = ptb.as_tensor_variable(np.asarray(sm0, dtype=sm1.dtype))
+            else:
+                sm0 = ptb.as_tensor_variable(sm0)
             if pytensor.scalar.upcast(sm0.dtype, sm1.dtype) == sm1.dtype:
                 lst2.append((ptb.cast(sm0, sm1.dtype), sM[1]))
 

tests/benchmarks/test_scan.py

@@ -449,7 +449,7 @@ def _test_scan_hessian_benchmark(mode, benchmark):
 
     def loss_outer(sum_outer, W):
         def loss_inner(sum_inner, W):
-            return sum_inner + (W**2).sum()
+            return sum_inner + (W ** np.int8(2)).sum()
 
         result_inner = scan(
             fn=loss_inner,
@@ -651,7 +651,7 @@ def cycle_step(A0, A1, A2, A1_hat, _norm, step_num):
 
             A0_L1_norm = linalg.norm(A0, ord=1)
 
-            return A0, A1, A2, A1_hat, A0_L1_norm, step_num + 1
+            return A0, A1, A2, A1_hat, A0_L1_norm, step_num + np.int8(1)
 
         return ifelse(
             norm < tol,

tests/link/numba/test_random.py

@@ -819,7 +819,7 @@ def test_rv_float_params_cast_to_float64(floatX):
     with pytensor.config.change_flags(floatX=floatX):
         rng = shared(np.random.default_rng(123))
 
-        # Opt-in: the int8 literals 0, 1 become float64 regardless of floatX.
+        # Opt-in: the integer literals 0, 1 become float64 regardless of floatX.
         assert _compiled_rv_dist_param_dtypes(
             pt.random.normal(0, 1, size=(5,), rng=rng)
         ) == ["float64", "float64"]
@@ -847,9 +847,10 @@ def test_rv_float_params_cast_to_float64(floatX):
 
 def test_rv_float_params_cast_respects_warn_float64():
     # When the user asked to be warned/raised on float64, the rewrite must not silently
-    # introduce it; the int8 params are left as-is.
+    # introduce it; the integer params are left as-is.
+    int_dtype = "int8" if pytensor.config.cast_policy == "custom" else "int64"
     with pytensor.config.change_flags(floatX="float32", warn_float64="raise"):
         rng = shared(np.random.default_rng(123))
         assert _compiled_rv_dist_param_dtypes(
             pt.random.normal(0, 1, size=(5,), rng=rng)
-        ) == ["int8", "int8"]
+        ) == [int_dtype, int_dtype]

tests/link/pytorch/test_basic.py

@@ -426,7 +426,7 @@ def test_ScalarLoop_while():
 def test_ScalarLoop_Elemwise_single_carries():
     n_steps = int64("n_steps")
     x0 = float64("x0")
-    x = x0 * 2
+    x = x0 * np.int8(2)
     until = x >= 10
 
     scalarop = ScalarLoop(init=[x0], update=[x], until=until)
@@ -452,8 +452,8 @@ def test_ScalarLoop_Elemwise_multi_carries():
     n_steps = int64("n_steps")
     x0 = float64("x0")
     x1 = float64("x1")
-    x = x0 * 2
-    x1_n = x1 * 3
+    x = x0 * np.int8(2)
+    x1_n = x1 * np.int8(3)
     until = x >= 10
 
     scalarop = ScalarLoop(init=[x0, x1], update=[x, x1_n], until=until)

tests/scalar/test_basic.py

@@ -467,7 +467,8 @@ def test_grad_abs():
 def test_constant():
     c = constant(2, name="a")
     assert c.name == "a"
-    assert c.dtype == "int8"
+    expected = "int8" if pytensor.config.cast_policy == "custom" else "int64"
+    assert c.dtype == expected
     c = constant(2, dtype="float32")
     assert c.name is None
     assert c.dtype == "float32"

tests/scalar/test_loop.py

@@ -221,7 +221,7 @@ def test_inner_composite(mode):
     n_steps = int64("n_steps")
     x = float64("x")
 
-    one = Composite([x], [cos(exp(x)) ** 2 + sin(exp(x)) ** 2])(x)
+    one = Composite([x], [cos(exp(x)) ** np.int8(2) + sin(exp(x)) ** np.int8(2)])(x)
 
     op = ScalarLoop(init=[x], update=[one + x])
     y = op(n_steps, x)
@@ -253,7 +253,7 @@ def test_inner_loop(mode):
     x = float64("x")
 
     x_in = float64("x_in")
-    inner_loop_op = ScalarLoop(init=[x_in], update=[x_in + 1])
+    inner_loop_op = ScalarLoop(init=[x_in], update=[x_in + np.int8(1)])
 
     outer_loop_op = ScalarLoop(
         init=[x], update=[inner_loop_op(n_steps, x)], constant=[n_steps]

tests/scan/test_basic.py

@@ -402,7 +402,7 @@ def test_no_step(self, mode, x_init):
         """We expect an empty output array when ``n_steps == 0``."""
 
         def f_pow(x_tm1):
-            return 2 * x_tm1
+            return np.int8(2) * x_tm1
 
         n_steps = iscalar("n_steps")
         values = scan(
@@ -437,7 +437,7 @@ def test_no_steps_sit_sot(self, mode, x, x_init):
         """We expect an empty output array when scanning over an empty sequence."""
 
         def inner_fn(x_seq, x_i):
-            return 2 * x_i
+            return np.int8(2) * x_i
 
         with config.change_flags(mode=mode):
             values = scan(

tests/tensor/rewriting/test_elemwise.py

@@ -479,47 +479,47 @@ def test_expansion_order(self):
             ),  # 15
             # test with constant
             (
-                (fw + fx) + (fy + fz) + 2.0,
+                (fw + fx) + (fy + fz) + np.float32(2.0),
                 (fw, fx, fy, fz),
                 (fwv, fxv, fyv, fzv),
                 1,
                 fwv + fxv + fyv + fzv + 2,
                 "float32",
             ),
             (
-                ((fw + fx) + 2.0 + fy) + fz,
+                ((fw + fx) + np.float32(2.0) + fy) + fz,
                 (fw, fx, fy, fz),
                 (fwv, fxv, fyv, fzv),
                 1,
                 fwv + fxv + fyv + fzv + 2,
                 "float32",
             ),
             (
-                (fw + (fx + 2.0 + fy)) + fz,
+                (fw + (fx + np.float32(2.0) + fy)) + fz,
                 (fw, fx, fy, fz),
                 (fwv, fxv, fyv, fzv),
                 1,
                 fwv + fxv + fyv + fzv + 2,
                 "float32",
             ),
             (
-                (fw + (fx + fy) + 2 + fz),
+                (fw + (fx + fy) + np.float32(2.0) + fz),
                 (fw, fx, fy, fz),
                 (fwv, fxv, fyv, fzv),
                 1,
                 fwv + fxv + fyv + fzv + 2,
                 "float32",
             ),
             (
-                fw + (fx + (fy + fz) + 2.0),
+                fw + (fx + (fy + fz) + np.float32(2.0)),
                 (fw, fx, fy, fz),
                 (fwv, fxv, fyv, fzv),
                 1,
                 fwv + fxv + fyv + fzv + 2,
                 "float32",
             ),  # 20
             (
-                2 + (fw + fx) + (fy + fz),
+                np.float32(2.0) + (fw + fx) + (fy + fz),
                 (fw, fx, fy, fz),
                 (fwv, fxv, fyv, fzv),
                 1,
@@ -658,7 +658,7 @@ def test_expansion_order(self):
                 "float32",
             ),
             (
-                fx - true_div(fy, 2),
+                fx - true_div(fy, np.float32(2.0)),
                 (fx, fy),
                 (fxv, fyv),
                 1,
@@ -685,7 +685,14 @@ def test_expansion_order(self):
                     "numpy": "float64",
                 },
             ),  # 40
-            (fx - (fy / 2), (fx, fy), (fxv, fyv), 1, fxv - (fyv / 2), "float32"),
+            (
+                fx - (fy / np.float32(2.0)),
+                (fx, fy),
+                (fxv, fyv),
+                1,
+                fxv - (fyv / 2),
+                "float32",
+            ),
             (
                 fx - (fy % fz),
                 (fx, fy, fz),
@@ -964,7 +971,10 @@ def test_expansion_order(self):
                     np.sum(-((fxv - fyv) ** 2) / 2),
                     -(fxv - fyv),
                 ),
-                ("float32", "float32"),
+                {
+                    "custom": ("float32", "float32"),
+                    "numpy+floatX": (config.floatX, "float32"),
+                },
             ),
             # Two Composite graphs that share the same input, but are split by
             # a non-elemwise operation (Assert)
@@ -1002,7 +1012,10 @@ def test_expansion_order(self):
                 (fxv,),
                 4,
                 (np.sum(fxv + 5) * np.exp(fxv) / (fxv + 5),),
-                ("float32",),
+                {
+                    "custom": ("float32",),
+                    "numpy+floatX": (config.floatX,),
+                },
             ),
             (
                 (sin(exp(fx)), exp(sin(fx))),

tests/tensor/rewriting/test_math.py

@@ -197,6 +197,11 @@ def inputs(xbc=(0, 0), ybc=(0, 0), zbc=(0, 0)):
     return x, y, z
 
 
+def _autocast_int_dtype():
+    """dtype the active cast_policy gives a small Python int literal."""
+    return "int8" if config.cast_policy == "custom" else "int64"
+
+
 def test_add_canonizer_problem0():
     n_segments = 10
     label = lscalar("label")
@@ -1507,7 +1512,6 @@ def test_or(self):
         for dtype, zero, one in [
             ("bool", np.array(False), np.array(True)),
             ("int8", np.int8(0), np.int8(1)),
-            ("int8", 0, 1),
         ]:
             x = scalar("x", dtype=dtype)
 
@@ -1986,7 +1990,7 @@ def test_log1pexp_log(self):
             f.maker.fgraph.outputs,
             [
                 pt.switch(
-                    x >= np.array([[0]], dtype=np.int8),
+                    x >= np.array([[0]], dtype=_autocast_int_dtype()),
                     pt.log1p(x),
                     np.array([[np.nan]], dtype=np.float32),
                 )
@@ -2011,7 +2015,7 @@ def test_log1mexp_log(self):
             f.maker.fgraph.outputs,
             [
                 pt.switch(
-                    x >= np.array([[0]], dtype=np.int8),
+                    x >= np.array([[0]], dtype=_autocast_int_dtype()),
                     pt.log1p(-x),
                     np.array([[np.nan]], dtype=np.float32),
                 )
@@ -2034,7 +2038,7 @@ def test_log1mexp_log1mexp(self):
             f.maker.fgraph.outputs,
             [
                 pt.switch(
-                    x <= np.array([[0]], dtype=np.int8),
+                    x <= np.array([[0]], dtype=_autocast_int_dtype()),
                     x,
                     np.array([[np.nan]], dtype=np.float32),
                 )
@@ -2093,7 +2097,7 @@ def test_sqr_sqrt(self):
         out = rewrite_graph(out, include=["canonicalize", "specialize", "stabilize"])
 
         expected = switch(
-            ge(x, np.zeros((1, 1), dtype="int8")),
+            ge(x, np.zeros((1, 1), dtype=_autocast_int_dtype())),
             x,
             np.full((1, 1), np.nan, dtype=out.type.dtype),
         )
@@ -2115,7 +2119,7 @@ def test_sqr_sqrt_integer_upcast(self):
         out = rewrite_graph(out, include=["canonicalize", "specialize", "stabilize"])
 
         expected = switch(
-            ge(x, np.zeros((1,), dtype="int8")),
+            ge(x, np.zeros((1,), dtype=_autocast_int_dtype())),
             x,
             np.full((1,), np.nan, dtype=dtype),
         )
@@ -4340,14 +4344,20 @@ def test_local_1msigmoid(self):
         xd = dscalar()
 
         # Test `exp_over_1_plus_exp`
-        f = pytensor.function([x], 1 - exp(x) / (1 + exp(x)), mode=m)
+        f = pytensor.function(
+            [x], np.float32(1) - exp(x) / (np.float32(1) + exp(x)), mode=m
+        )
         # FIXME: PatternNodeRewriter does not copy stack trace
         #  (see https://github.com/Theano/Theano/issues/4581)
         # assert check_stack_trace(f, ops_to_check=[neg, sigmoid])
         assert equal_computations(f.maker.fgraph.outputs, [sigmoid(-x)])
 
         # Test `inv_1_plus_exp`
-        f = pytensor.function([x], 1 - pt.fill(x, 1.0) / (1 + exp(-x)), mode=m)
+        f = pytensor.function(
+            [x],
+            np.float32(1) - pt.fill(x, np.float32(1.0)) / (np.float32(1) + exp(-x)),
+            mode=m,
+        )
         # assert check_stack_trace(f, ops_to_check=[neg, sigmoid])
         assert equal_computations(f.maker.fgraph.outputs, [sigmoid(-x)])
 
@@ -4622,7 +4632,7 @@ def test_local_logit_sigmoid():
     """Test that graphs of the form ``logit(sigmoid(x))`` and ``sigmoid(logit(x))`` get rewritten to ``x``."""
 
     def logit_fn(x):
-        return log(x / (1 - x))
+        return log(x / (np.float32(1) - x))
 
     x = fmatrix()
 

tests/tensor/test_blas.py

@@ -1278,15 +1278,17 @@ def test_gemv1(self):
         self.t_gemv1((3, 0))
         self.t_gemv1((0, 0))
 
-    def test_gemv2(self):
+    @pytest.mark.parametrize("cast_policy", ["custom", "numpy+floatX"])
+    def test_gemv2(self, cast_policy):
         # test vector2+dot(vector1,matrix)
         rng = np.random.default_rng(unittest_tools.fetch_seed())
         v1 = shared(np.array(rng.uniform(size=(2,)), dtype="float32"))
         v2_orig = np.array(rng.uniform(size=(3,)), dtype="float32")
         v2 = shared(v2_orig)
         m = shared(np.array(rng.uniform(size=(2, 3)), dtype="float32"))
 
-        f = function([], v2 + pytensor.tensor.dot(v1, m), mode=mode_blas_opt)
+        with config.change_flags(cast_policy=cast_policy):
+            f = function([], v2 + pytensor.tensor.dot(v1, m), mode=mode_blas_opt)
 
         # Assert they produce the same output
         assert np.allclose(f(), np.dot(v1.get_value(), m.get_value()) + v2.get_value())
@@ -1295,9 +1297,13 @@ def test_gemv2(self):
         assert topo[-1].op.inplace is False
 
         # test the inplace version
-        g = function(
-            [], [], updates=[(v2, v2 + pytensor.tensor.dot(v1, m))], mode=mode_blas_opt
-        )
+        with config.change_flags(cast_policy=cast_policy):
+            g = function(
+                [],
+                [],
+                updates=[(v2, v2 + pytensor.tensor.dot(v1, m))],
+                mode=mode_blas_opt,
+            )
 
         # Assert they produce the same output
         g()

tests/test_gradient.py

@@ -570,7 +570,7 @@ def test_downcast_dtype(self):
         # get upcasted to float64.
         # x has dtype float32, regardless of the value of floatX
         x = fscalar("x")
-        y = x * 2
+        y = x * np.float32(2)
         z = lscalar("z")
 
         c = y + z