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/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/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, 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/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))), 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() 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() 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