From ba25d862d6a015266a0846d5dd0f82bb7ac2877e Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 7 Jul 2026 07:54:32 +0000 Subject: [PATCH 1/5] Fix gemm/gemv fusion dropped by literal scale upcast Under cast_policy="numpy+floatX" the +/-1.0 scale seeded by _gemm_canonicalize autocasts to float64, fails the upcast check in _gemm_from_factored_list, and the whole gemm/gemv candidate is silently discarded (e.g. v2 + dot(v1, m) keeps a separate Add instead of folding beta into Gemv). Build literal scales with the matrix dtype instead. No-op under the default custom policy; test_gemv2 is parametrized on cast_policy to cover both. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5 --- pytensor/tensor/rewriting/blas.py | 8 +++++++- tests/tensor/test_blas.py | 16 +++++++++++----- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/pytensor/tensor/rewriting/blas.py b/pytensor/tensor/rewriting/blas.py index 03a1e8b0ab..2cd1cd3467 100644 --- a/pytensor/tensor/rewriting/blas.py +++ b/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])) diff --git a/tests/tensor/test_blas.py b/tests/tensor/test_blas.py index b9ebc5939c..973f19650d 100644 --- a/tests/tensor/test_blas.py +++ b/tests/tensor/test_blas.py @@ -1278,7 +1278,8 @@ 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")) @@ -1286,7 +1287,8 @@ def test_gemv2(self): 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() From f4985c4b3342275500f97db6173f6bfee0e69e2c Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 7 Jul 2026 07:54:40 +0000 Subject: [PATCH 2/5] Make rewriting test expectations cast_policy-aware Add a _autocast_int_dtype() helper and use it for the TestExpLog/ TestSqrSqrt switch-constant expectations, drop the now-obsolete ("int8", 0, 1) case from the bitwise-or test, and pin the float32 literals in the two sigmoid rewrite tests. Green under both policies. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5 --- tests/tensor/rewriting/test_math.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py index 8c2192096a..3b117107e7 100644 --- a/tests/tensor/rewriting/test_math.py +++ b/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() From 1f0050f567471f8fea8c019080a7381cf076b6ea Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 7 Jul 2026 07:54:45 +0000 Subject: [PATCH 3/5] Make fusion tests cast_policy-aware Pin the +2.0 and /2 fusion constants to np.float32 so the fused output stays float32, and assert the honest per-policy out_dtype for the multi-output cases 72/74 via the {"custom": ..., "numpy+floatX": ...} convention. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5 --- tests/tensor/rewriting/test_elemwise.py | 33 +++++++++++++++++-------- 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/tests/tensor/rewriting/test_elemwise.py b/tests/tensor/rewriting/test_elemwise.py index 74de45282e..d453b3f313 100644 --- a/tests/tensor/rewriting/test_elemwise.py +++ b/tests/tensor/rewriting/test_elemwise.py @@ -479,7 +479,7 @@ 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, @@ -487,7 +487,7 @@ def test_expansion_order(self): "float32", ), ( - ((fw + fx) + 2.0 + fy) + fz, + ((fw + fx) + np.float32(2.0) + fy) + fz, (fw, fx, fy, fz), (fwv, fxv, fyv, fzv), 1, @@ -495,7 +495,7 @@ def test_expansion_order(self): "float32", ), ( - (fw + (fx + 2.0 + fy)) + fz, + (fw + (fx + np.float32(2.0) + fy)) + fz, (fw, fx, fy, fz), (fwv, fxv, fyv, fzv), 1, @@ -503,7 +503,7 @@ def test_expansion_order(self): "float32", ), ( - (fw + (fx + fy) + 2 + fz), + (fw + (fx + fy) + np.float32(2.0) + fz), (fw, fx, fy, fz), (fwv, fxv, fyv, fzv), 1, @@ -511,7 +511,7 @@ def test_expansion_order(self): "float32", ), ( - fw + (fx + (fy + fz) + 2.0), + fw + (fx + (fy + fz) + np.float32(2.0)), (fw, fx, fy, fz), (fwv, fxv, fyv, fzv), 1, @@ -519,7 +519,7 @@ def test_expansion_order(self): "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))), From 31321595437854068fb57ab14df09e6fc324608d Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 7 Jul 2026 07:54:51 +0000 Subject: [PATCH 4/5] Make scan, scalar, gradient, and backend tests cast_policy-aware Pin the inner-graph integer literals that would otherwise promote under numpy+floatX (scan 2*x, scalar loop ** 2 / + 1, gradient x * 2, pytorch Elemwise-ScalarLoop x * 2 / x1 * 3), and make the scalar test_constant and numba float-param dtype expectations derive from the active cast_policy. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5 --- tests/link/numba/test_random.py | 7 ++++--- tests/link/pytorch/test_basic.py | 6 +++--- tests/scalar/test_basic.py | 3 ++- tests/scalar/test_loop.py | 4 ++-- tests/scan/test_basic.py | 4 ++-- tests/test_gradient.py | 2 +- 6 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/link/numba/test_random.py b/tests/link/numba/test_random.py index 252967a0e8..7ee4f7523b 100644 --- a/tests/link/numba/test_random.py +++ b/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] diff --git a/tests/link/pytorch/test_basic.py b/tests/link/pytorch/test_basic.py index 6ec9f34218..0414065566 100644 --- a/tests/link/pytorch/test_basic.py +++ b/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) diff --git a/tests/scalar/test_basic.py b/tests/scalar/test_basic.py index 73a074e73e..334fe75a71 100644 --- a/tests/scalar/test_basic.py +++ b/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" diff --git a/tests/scalar/test_loop.py b/tests/scalar/test_loop.py index 6cca6cbaf9..6b8bed78da 100644 --- a/tests/scalar/test_loop.py +++ b/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] diff --git a/tests/scan/test_basic.py b/tests/scan/test_basic.py index 1f7a5fdc73..e13a9d161c 100644 --- a/tests/scan/test_basic.py +++ b/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( diff --git a/tests/test_gradient.py b/tests/test_gradient.py index c4ad6b3b22..122905c0b9 100644 --- a/tests/test_gradient.py +++ b/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 From 3f5795da94a8695b5327e52836e0e85c854b5c92 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 7 Jul 2026 11:19:35 +0000 Subject: [PATCH 5/5] Make scan benchmarks and autocast doctest cast_policy-aware The hessian benchmark's inner W**2 and cyclical reduction's step_num + 1 promote under numpy+floatX (float32 -> float64 and int32 -> int64), so scan rejects the state dtype and IfElse rejects the branches; pin both literals to int8. The autocast_float_as doctest only makes sense under custom, so scope it with change_flags. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5 --- pytensor/scalar/basic.py | 8 +++++--- tests/benchmarks/test_scan.py | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py index ae03530fa6..6491a83bad 100644 --- a/pytensor/scalar/basic.py +++ b/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 """ diff --git a/tests/benchmarks/test_scan.py b/tests/benchmarks/test_scan.py index 3faf11b134..a9d31772fd 100644 --- a/tests/benchmarks/test_scan.py +++ b/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,