From 639ef8c653e663d8f1f2a2b788193968b2955e52 Mon Sep 17 00:00:00 2001 From: Trey Wenger Date: Fri, 1 Mar 2024 18:23:48 -0600 Subject: [PATCH 1/6] default acc_dtype to floatX --- pytensor/tensor/elemwise.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pytensor/tensor/elemwise.py b/pytensor/tensor/elemwise.py index fce2be9d2a..0329575c01 100644 --- a/pytensor/tensor/elemwise.py +++ b/pytensor/tensor/elemwise.py @@ -1246,10 +1246,9 @@ def __init__( The dtype of the internal accumulator. If ``None`` (default), we use the dtype in the list below, or the input dtype if its precision is higher: - - for int dtypes, we use at least int64; - for uint dtypes, we use at least uint64; - - for float dtypes, we use at least float64; + - for float dtypes, we use at least floatX; - for complex dtypes, we use at least complex128. upcast_discrete_output See @@ -1333,8 +1332,8 @@ def _acc_dtype(self, idtype): uint8="uint64", uint16="uint64", uint32="uint64", - float16="float32", - float32="float64", + float16=config.floatX, + float32=config.floatX, complex64="complex128", ).get(idtype, idtype) elif acc_dtype in continuous_dtypes and idtype in discrete_dtypes: From da561b0f1bb814462daf13343af9c4d4e0b17fb3 Mon Sep 17 00:00:00 2001 From: Trey Wenger Date: Fri, 1 Mar 2024 18:38:18 -0600 Subject: [PATCH 2/6] update tests for default acc_dtype --- tests/tensor/test_math.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py index e346348406..0ce8dd252f 100644 --- a/tests/tensor/test_math.py +++ b/tests/tensor/test_math.py @@ -3030,8 +3030,8 @@ def test_reduce_default_acc_dtype(self): uint8="uint64", uint16="uint64", uint32="uint64", - float16="float32", - float32="float64", + float16=config.floatX, + float32=config.floatX, complex64="complex128", ).get(dtype, dtype) f = function([x], s, mode=self.mode) @@ -3255,8 +3255,8 @@ def test_prod_without_zeros_default_acc_dtype(self): uint8="uint64", uint16="uint64", uint32="uint64", - float16="float32", - float32="float64", + float16=config.floatX, + float32=config.floatX, complex64="complex128", ).get(dtype, dtype) From 1494fa0a1f83a4baa71656e5a3fdde79b915eb13 Mon Sep 17 00:00:00 2001 From: Trey Wenger Date: Sat, 2 Mar 2024 12:37:03 -0600 Subject: [PATCH 3/6] enforce local_careduce_fusion acc_dtype --- pytensor/tensor/rewriting/elemwise.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/elemwise.py b/pytensor/tensor/rewriting/elemwise.py index 6453aecb94..74a7a0cdbd 100644 --- a/pytensor/tensor/rewriting/elemwise.py +++ b/pytensor/tensor/rewriting/elemwise.py @@ -1176,7 +1176,7 @@ def local_careduce_fusion(fgraph, node): car_acc_dtype = node.op.acc_dtype scalar_elm_inputs = [ - ps.get_scalar_type(inp.type.dtype).make_variable() for inp in elm_inputs + ps.get_scalar_type(car_acc_dtype).make_variable() for inp in elm_inputs ] elm_output = elm_scalar_op(*scalar_elm_inputs) From 7dfe4525730cd4f69d5618622439a87d80a09dac Mon Sep 17 00:00:00 2001 From: Trey Wenger Date: Sat, 2 Mar 2024 14:46:03 -0600 Subject: [PATCH 4/6] revert CAReduce acc_dtype, ensure verify_grad is casting random_projection to output dtype --- pytensor/tensor/rewriting/elemwise.py | 2 +- tests/tensor/test_extra_ops.py | 4 +++- tests/tensor/test_subtensor.py | 14 ++++++++++++-- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/pytensor/tensor/rewriting/elemwise.py b/pytensor/tensor/rewriting/elemwise.py index 74a7a0cdbd..6453aecb94 100644 --- a/pytensor/tensor/rewriting/elemwise.py +++ b/pytensor/tensor/rewriting/elemwise.py @@ -1176,7 +1176,7 @@ def local_careduce_fusion(fgraph, node): car_acc_dtype = node.op.acc_dtype scalar_elm_inputs = [ - ps.get_scalar_type(car_acc_dtype).make_variable() for inp in elm_inputs + ps.get_scalar_type(inp.type.dtype).make_variable() for inp in elm_inputs ] elm_output = elm_scalar_op(*scalar_elm_inputs) diff --git a/tests/tensor/test_extra_ops.py b/tests/tensor/test_extra_ops.py index cda745d023..c948ec9679 100644 --- a/tests/tensor/test_extra_ops.py +++ b/tests/tensor/test_extra_ops.py @@ -631,7 +631,9 @@ def test_grad(self, ndim): a = np.random.random((10,) * ndim).astype(config.floatX) for axis in self._possible_axis(ndim): - utt.verify_grad(lambda x: Repeat(axis=axis)(x, 3), [a]) + utt.verify_grad( + lambda x: Repeat(axis=axis)(x, 3), [a], cast_to_output_type=True + ) def test_broadcastable(self): x = TensorType(config.floatX, shape=(None, 1, None))() diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py index 946cb48b6b..d01bfd39db 100644 --- a/tests/tensor/test_subtensor.py +++ b/tests/tensor/test_subtensor.py @@ -771,8 +771,18 @@ def test_grad_2d_inc_set_subtensor(self): t = op(n[:z, :z], m) gn, gm = pytensor.grad(pt_sum(t), [n, m]) - utt.verify_grad(lambda m: op(n[:z, :z], m), [mv], mode=self.mode) - utt.verify_grad(lambda nn: op(nn[:z, :z], mv), [data], mode=self.mode) + utt.verify_grad( + lambda m: op(n[:z, :z], m), + [mv], + mode=self.mode, + cast_to_output_type=True, + ) + utt.verify_grad( + lambda nn: op(nn[:z, :z], mv), + [data], + mode=self.mode, + cast_to_output_type=True, + ) def test_grad_0d(self): data = np.asarray(random(2, 3), dtype=self.dtype) From 6c931482f42a9e41c8484113538f3fa60ac5874c Mon Sep 17 00:00:00 2001 From: Trey Wenger Date: Sat, 2 Mar 2024 15:12:13 -0600 Subject: [PATCH 5/6] revert test cast_to_output_type, change default to True in verify_grad --- pytensor/gradient.py | 2 +- tests/tensor/test_extra_ops.py | 4 +--- tests/tensor/test_subtensor.py | 2 -- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/pytensor/gradient.py b/pytensor/gradient.py index 0d18070dcc..c9007e05f1 100644 --- a/pytensor/gradient.py +++ b/pytensor/gradient.py @@ -1673,7 +1673,7 @@ def verify_grad( abs_tol: Optional[float] = None, rel_tol: Optional[float] = None, mode: Optional[Union["Mode", str]] = None, - cast_to_output_type: bool = False, + cast_to_output_type: bool = True, no_debug_ref: bool = True, ): """Test a gradient by Finite Difference Method. Raise error on failure. diff --git a/tests/tensor/test_extra_ops.py b/tests/tensor/test_extra_ops.py index c948ec9679..cda745d023 100644 --- a/tests/tensor/test_extra_ops.py +++ b/tests/tensor/test_extra_ops.py @@ -631,9 +631,7 @@ def test_grad(self, ndim): a = np.random.random((10,) * ndim).astype(config.floatX) for axis in self._possible_axis(ndim): - utt.verify_grad( - lambda x: Repeat(axis=axis)(x, 3), [a], cast_to_output_type=True - ) + utt.verify_grad(lambda x: Repeat(axis=axis)(x, 3), [a]) def test_broadcastable(self): x = TensorType(config.floatX, shape=(None, 1, None))() diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py index d01bfd39db..e692939b18 100644 --- a/tests/tensor/test_subtensor.py +++ b/tests/tensor/test_subtensor.py @@ -775,13 +775,11 @@ def test_grad_2d_inc_set_subtensor(self): lambda m: op(n[:z, :z], m), [mv], mode=self.mode, - cast_to_output_type=True, ) utt.verify_grad( lambda nn: op(nn[:z, :z], mv), [data], mode=self.mode, - cast_to_output_type=True, ) def test_grad_0d(self): From d8e5de52a1981d20edaf66e1f06daa273432bd8e Mon Sep 17 00:00:00 2001 From: Trey Wenger Date: Sat, 2 Mar 2024 17:21:52 -0600 Subject: [PATCH 6/6] revert verify_grad default, enforce cast_to_output_type for offending tests --- pytensor/gradient.py | 2 +- tests/tensor/test_extra_ops.py | 4 +++- tests/tensor/test_subtensor.py | 2 ++ tests/tensor/utils.py | 1 + 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/pytensor/gradient.py b/pytensor/gradient.py index c9007e05f1..0d18070dcc 100644 --- a/pytensor/gradient.py +++ b/pytensor/gradient.py @@ -1673,7 +1673,7 @@ def verify_grad( abs_tol: Optional[float] = None, rel_tol: Optional[float] = None, mode: Optional[Union["Mode", str]] = None, - cast_to_output_type: bool = True, + cast_to_output_type: bool = False, no_debug_ref: bool = True, ): """Test a gradient by Finite Difference Method. Raise error on failure. diff --git a/tests/tensor/test_extra_ops.py b/tests/tensor/test_extra_ops.py index cda745d023..c948ec9679 100644 --- a/tests/tensor/test_extra_ops.py +++ b/tests/tensor/test_extra_ops.py @@ -631,7 +631,9 @@ def test_grad(self, ndim): a = np.random.random((10,) * ndim).astype(config.floatX) for axis in self._possible_axis(ndim): - utt.verify_grad(lambda x: Repeat(axis=axis)(x, 3), [a]) + utt.verify_grad( + lambda x: Repeat(axis=axis)(x, 3), [a], cast_to_output_type=True + ) def test_broadcastable(self): x = TensorType(config.floatX, shape=(None, 1, None))() diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py index e692939b18..d01bfd39db 100644 --- a/tests/tensor/test_subtensor.py +++ b/tests/tensor/test_subtensor.py @@ -775,11 +775,13 @@ def test_grad_2d_inc_set_subtensor(self): lambda m: op(n[:z, :z], m), [mv], mode=self.mode, + cast_to_output_type=True, ) utt.verify_grad( lambda nn: op(nn[:z, :z], mv), [data], mode=self.mode, + cast_to_output_type=True, ) def test_grad_0d(self): diff --git a/tests/tensor/utils.py b/tests/tensor/utils.py index be5f2a029e..d2a556482f 100644 --- a/tests/tensor/utils.py +++ b/tests/tensor/utils.py @@ -592,6 +592,7 @@ def test_grad(self): mode=self.mode, rel_tol=_grad_rtol, eps=_grad_eps, + cast_to_output_type=True, ) except Exception as exc: err_msg = (