fix the bug

Author: ferrine

pymc/logprob/transforms.py

@@ -40,6 +40,7 @@
 from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import pytensor
 import pytensor.tensor as pt
 
 from pytensor import scan
@@ -959,7 +960,8 @@ class SimplexTransform(RVTransform):
 
     def forward(self, value, *inputs):
         log_value = pt.log(value)
-        shift = pt.sum(log_value, -1, keepdims=True) / value.shape[-1]
+        N = value.shape[-1].astype(pytensor.config.floatX)
+        shift = pt.sum(log_value, -1, keepdims=True) / N
         return log_value[..., :-1] - shift
 
     def backward(self, value, *inputs):
@@ -969,6 +971,7 @@ def backward(self, value, *inputs):
 
     def log_jac_det(self, value, *inputs):
         N = value.shape[-1] + 1
+        N = N.astype(pytensor.config.floatX)
         sum_value = pt.sum(value, -1, keepdims=True)
         value_sum_expanded = value + sum_value
         value_sum_expanded = pt.concatenate([value_sum_expanded, pt.zeros(sum_value.shape)], -1)

tests/distributions/test_multivariate.py

@@ -1716,13 +1716,13 @@ class TestDirichlet(BaseTestDistributionRandom):
 
     @pytensor.config.change_flags(warn_float64="raise", floatX="float32")
     def test_dirichlet_float32(self):
-        """https://github.com/pymc-devs/pymc/issues/6779
-        """
+        """https://github.com/pymc-devs/pymc/issues/6779"""
         with pm.Model() as model:
             c = pm.floatX([1, 1, 1])
             pm.Dirichlet("a", c)
         model.point_logps()
 
+
 class TestMultinomial(BaseTestDistributionRandom):
     pymc_dist = pm.Multinomial
     pymc_dist_params = {"n": 85, "p": np.array([0.28, 0.62, 0.10])}
@@ -2134,4 +2134,3 @@ def test_posdef_symmetric(matrix, result):
     """
     data = np.array(matrix, dtype=pytensor.config.floatX)
     assert posdef(data) == result
-