Parametrize some more tests by floatX

Author: michaelosthege

pymc/distributions/simulator.py

@@ -295,11 +295,9 @@ def __call__(self, epsilon, obs_data, sim_data):
         return self.d_n * np.sum(-np.log(nu_d / self.rho_d) / epsilon) + self.log_r
 
 
-scalarX = at.dscalar if aesara.config.floatX == "float64" else at.fscalar
-vectorX = at.dvector if aesara.config.floatX == "float64" else at.fvector
-
-
 def create_sum_stat_op_from_fn(fn):
+    vectorX = at.dvector if aesara.config.floatX == "float64" else at.fvector
+
     # Check if callable returns TensorVariable with dummy inputs
     try:
         res = fn(vectorX())
@@ -322,6 +320,9 @@ def perform(self, node, inputs, outputs):
 
 
 def create_distance_op_from_fn(fn):
+    scalarX = at.dscalar if aesara.config.floatX == "float64" else at.fscalar
+    vectorX = at.dvector if aesara.config.floatX == "float64" else at.fvector
+
     # Check if callable returns TensorVariable with dummy inputs
     try:
         res = fn(scalarX(), vectorX(), vectorX())

pymc/tests/test_aesaraf.py

@@ -47,9 +47,6 @@
 from pymc.exceptions import ShapeError
 from pymc.vartypes import int_types
 
-FLOATX = str(aesara.config.floatX)
-INTX = str(_conversion_map[FLOATX])
-
 
 def test_change_rv_size():
     loc = at.as_tensor_variable([1, 2])
@@ -176,57 +173,59 @@ def setup_class(self):
         self.output_buffer = dict()
         self.func_buffer = dict()
 
-    def _input_tensors(self, shape):
+    def _input_tensors(self, shape, floatX):
+        intX = str(_conversion_map[floatX])
         ndim = len(shape)
-        arr = TensorType(FLOATX, [False] * ndim)("arr")
-        indices = TensorType(INTX, [False] * ndim)("indices")
-        arr.tag.test_value = np.zeros(shape, dtype=FLOATX)
-        indices.tag.test_value = np.zeros(shape, dtype=INTX)
+        arr = TensorType(floatX, [False] * ndim)("arr")
+        indices = TensorType(intX, [False] * ndim)("indices")
+        arr.tag.test_value = np.zeros(shape, dtype=floatX)
+        indices.tag.test_value = np.zeros(shape, dtype=intX)
         return arr, indices
 
-    def get_input_tensors(self, shape):
+    def get_input_tensors(self, shape, floatX):
         ndim = len(shape)
         try:
-            return self.inputs_buffer[ndim]
+            return self.inputs_buffer[(ndim, floatX)]
         except KeyError:
-            arr, indices = self._input_tensors(shape)
-            self.inputs_buffer[ndim] = arr, indices
+            arr, indices = self._input_tensors(shape, floatX)
+            self.inputs_buffer[(ndim, floatX)] = arr, indices
             return arr, indices
 
     def _output_tensor(self, arr, indices, axis):
         return take_along_axis(arr, indices, axis)
 
-    def get_output_tensors(self, shape, axis):
+    def get_output_tensors(self, shape, axis, floatX):
         ndim = len(shape)
         try:
-            return self.output_buffer[(ndim, axis)]
+            return self.output_buffer[(ndim, axis, floatX)]
         except KeyError:
-            arr, indices = self.get_input_tensors(shape)
+            arr, indices = self.get_input_tensors(shape, floatX)
             out = self._output_tensor(arr, indices, axis)
-            self.output_buffer[(ndim, axis)] = out
+            self.output_buffer[(ndim, axis, floatX)] = out
             return out
 
     def _function(self, arr, indices, out):
         return aesara.function([arr, indices], [out])
 
-    def get_function(self, shape, axis):
+    def get_function(self, shape, axis, floatX):
         ndim = len(shape)
         try:
-            return self.func_buffer[(ndim, axis)]
+            return self.func_buffer[(ndim, axis, floatX)]
         except KeyError:
-            arr, indices = self.get_input_tensors(shape)
-            out = self.get_output_tensors(shape, axis)
+            arr, indices = self.get_input_tensors(shape, floatX)
+            out = self.get_output_tensors(shape, axis, floatX)
             func = self._function(arr, indices, out)
-            self.func_buffer[(ndim, axis)] = func
+            self.func_buffer[(ndim, axis, floatX)] = func
             return func
 
     @staticmethod
-    def get_input_values(shape, axis, samples):
-        arr = np.random.randn(*shape).astype(FLOATX)
+    def get_input_values(shape, axis, samples, floatX):
+        intX = str(_conversion_map[floatX])
+        arr = np.random.randn(*shape).astype(floatX)
         size = list(shape)
         size[axis] = samples
         size = tuple(size)
-        indices = np.random.randint(low=0, high=shape[axis], size=size, dtype=INTX)
+        indices = np.random.randint(low=0, high=shape[axis], size=size, dtype=intX)
         return arr, indices
 
     @pytest.mark.parametrize(
@@ -250,10 +249,12 @@ def get_input_values(shape, axis, samples):
         ),
         ids=str,
     )
-    def test_take_along_axis(self, shape, axis, samples):
-        arr, indices = self.get_input_values(shape, axis, samples)
-        func = self.get_function(shape, axis)
-        assert np.allclose(np_take_along_axis(arr, indices, axis=axis), func(arr, indices)[0])
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    def test_take_along_axis(self, shape, axis, samples, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            arr, indices = self.get_input_values(shape, axis, samples, floatX)
+            func = self.get_function(shape, axis, floatX)
+            assert np.allclose(np_take_along_axis(arr, indices, axis=axis), func(arr, indices)[0])
 
     @pytest.mark.parametrize(
         ["shape", "axis", "samples"],
@@ -276,53 +277,62 @@ def test_take_along_axis(self, shape, axis, samples):
         ),
         ids=str,
     )
-    def test_take_along_axis_grad(self, shape, axis, samples):
-        if axis < 0:
-            _axis = len(shape) + axis
-        else:
-            _axis = axis
-        # Setup the aesara function
-        t_arr, t_indices = self.get_input_tensors(shape)
-        t_out2 = aesara.grad(
-            at.sum(self._output_tensor(t_arr**2, t_indices, axis)),
-            t_arr,
-        )
-        func = aesara.function([t_arr, t_indices], [t_out2])
-
-        # Test that the gradient gives the same output as what is expected
-        arr, indices = self.get_input_values(shape, axis, samples)
-        expected_grad = np.zeros_like(arr)
-        slicer = [slice(None)] * len(shape)
-        for i in range(indices.shape[axis]):
-            slicer[axis] = i
-            inds = indices[tuple(slicer)].reshape(shape[:_axis] + (1,) + shape[_axis + 1 :])
-            inds = _make_along_axis_idx(shape, inds, _axis)
-            expected_grad[inds] += 1
-        expected_grad *= 2 * arr
-        out = func(arr, indices)[0]
-        assert np.allclose(out, expected_grad)
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    def test_take_along_axis_grad(self, shape, axis, samples, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            if axis < 0:
+                _axis = len(shape) + axis
+            else:
+                _axis = axis
+            # Setup the aesara function
+            t_arr, t_indices = self.get_input_tensors(shape, floatX)
+            t_out2 = aesara.grad(
+                at.sum(self._output_tensor(t_arr**2, t_indices, axis)),
+                t_arr,
+            )
+            func = aesara.function([t_arr, t_indices], [t_out2])
+
+            # Test that the gradient gives the same output as what is expected
+            arr, indices = self.get_input_values(shape, axis, samples, floatX)
+            expected_grad = np.zeros_like(arr)
+            slicer = [slice(None)] * len(shape)
+            for i in range(indices.shape[axis]):
+                slicer[axis] = i
+                inds = indices[tuple(slicer)].reshape(shape[:_axis] + (1,) + shape[_axis + 1 :])
+                inds = _make_along_axis_idx(shape, inds, _axis)
+                expected_grad[inds] += 1
+            expected_grad *= 2 * arr
+            out = func(arr, indices)[0]
+            assert np.allclose(out, expected_grad)
 
     @pytest.mark.parametrize("axis", [-4, 4], ids=str)
-    def test_axis_failure(self, axis):
-        arr, indices = self.get_input_tensors((3, 1))
-        with pytest.raises(ValueError):
-            take_along_axis(arr, indices, axis=axis)
-
-    def test_ndim_failure(self):
-        arr = TensorType(FLOATX, [False] * 3)("arr")
-        indices = TensorType(INTX, [False] * 2)("indices")
-        arr.tag.test_value = np.zeros((1,) * arr.ndim, dtype=FLOATX)
-        indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=INTX)
-        with pytest.raises(ValueError):
-            take_along_axis(arr, indices)
-
-    def test_dtype_failure(self):
-        arr = TensorType(FLOATX, [False] * 3)("arr")
-        indices = TensorType(FLOATX, [False] * 3)("indices")
-        arr.tag.test_value = np.zeros((1,) * arr.ndim, dtype=FLOATX)
-        indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=FLOATX)
-        with pytest.raises(IndexError):
-            take_along_axis(arr, indices)
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    def test_axis_failure(self, axis, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            arr, indices = self.get_input_tensors((3, 1), floatX)
+            with pytest.raises(ValueError):
+                take_along_axis(arr, indices, axis=axis)
+
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    def test_ndim_failure(self, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            intX = str(_conversion_map[floatX])
+            arr = TensorType(floatX, [False] * 3)("arr")
+            indices = TensorType(intX, [False] * 2)("indices")
+            arr.tag.test_value = np.zeros((1,) * arr.ndim, dtype=floatX)
+            indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=intX)
+            with pytest.raises(ValueError):
+                take_along_axis(arr, indices)
+
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    def test_dtype_failure(self, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            arr = TensorType(floatX, [False] * 3)("arr")
+            indices = TensorType(floatX, [False] * 3)("indices")
+            arr.tag.test_value = np.zeros((1,) * arr.ndim, dtype=floatX)
+            indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=floatX)
+            with pytest.raises(IndexError):
+                take_along_axis(arr, indices)
 
 
 def test_extract_obs_data():

pymc/tests/test_distributions_random.py

@@ -14,6 +14,7 @@
 import functools
 import itertools
 import re
+import sys
 
 from typing import Callable, List, Optional
 
@@ -1571,14 +1572,18 @@ def constant_rng_fn(self, size, c):
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
         "check_rv_size",
-        "check_dtype",
     ]
 
-    def check_dtype(self):
-        assert pm.Constant.dist(2**4).dtype == "int8"
-        assert pm.Constant.dist(2**16).dtype == "int32"
-        assert pm.Constant.dist(2**32).dtype == "int64"
-        assert pm.Constant.dist(2.0).dtype == aesara.config.floatX
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    @pytest.mark.xfail(
+        sys.platform == "win32", reason="https://github.com/aesara-devs/aesara/issues/871"
+    )
+    def test_dtype(self, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            assert pm.Constant.dist(2**4).dtype == "int8"
+            assert pm.Constant.dist(2**16).dtype == "int32"
+            assert pm.Constant.dist(2**32).dtype == "int64"
+            assert pm.Constant.dist(2.0).dtype == floatX
 
 
 class TestOrderedLogistic(BaseTestDistributionRandom):

pymc/tests/test_mixture.py

@@ -659,24 +659,26 @@ def test_iterable_single_component_warning(self):
         with pytest.warns(UserWarning, match="Single component will be treated as a mixture"):
             Mixture.dist(w=[0.5, 0.5], comp_dists=[Normal.dist(size=2)])
 
-    def test_mixture_dtype(self):
-        mix_dtype = Mixture.dist(
-            w=[0.5, 0.5],
-            comp_dists=[
-                Multinomial.dist(n=5, p=[0.5, 0.5]),
-                Multinomial.dist(n=5, p=[0.5, 0.5]),
-            ],
-        ).dtype
-        assert mix_dtype == "int64"
-
-        mix_dtype = Mixture.dist(
-            w=[0.5, 0.5],
-            comp_dists=[
-                Dirichlet.dist(a=[0.5, 0.5]),
-                Dirichlet.dist(a=[0.5, 0.5]),
-            ],
-        ).dtype
-        assert mix_dtype == aesara.config.floatX
+    @pytest.mark.parametrize("floatX", ["float32", "float64"])
+    def test_mixture_dtype(self, floatX):
+        with aesara.config.change_flags(floatX=floatX):
+            mix_dtype = Mixture.dist(
+                w=[0.5, 0.5],
+                comp_dists=[
+                    Multinomial.dist(n=5, p=[0.5, 0.5]),
+                    Multinomial.dist(n=5, p=[0.5, 0.5]),
+                ],
+            ).dtype
+            assert mix_dtype == "int64"
+
+            mix_dtype = Mixture.dist(
+                w=[0.5, 0.5],
+                comp_dists=[
+                    Dirichlet.dist(a=[0.5, 0.5]),
+                    Dirichlet.dist(a=[0.5, 0.5]),
+                ],
+            ).dtype
+            assert mix_dtype == floatX
 
     @pytest.mark.parametrize(
         "comp_dists, expected_shape",