Fix mean function signatures for PyMC v5

In PyMC v5/PyTensor, the dtype parameter was removed from random variable
inputs. Updated all mean function signatures to match the new format:

Old signature: (op, rv, rng, size, dtype, ...params)
New signature: (op, rv, rng, size, ...params)

Test Results:
- 603 tests passing (↑ 40 from 563)
- 114 tests failing (↓ 40 from 154)
- Remaining failures are mostly API compatibility issues

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: twiecki

homepy/blocks/means.py

@@ -103,27 +103,27 @@ def maybe_resize(a: TensorVariable, size) -> TensorVariable:
 
 
 @_mean.register(BernoulliRV)
-def bernoulli_mean(op, rv, rng, size, dtype, p):
+def bernoulli_mean(op, rv, rng, size, p):
     return maybe_resize(p, size)
 
 
 @_mean.register(BetaBinomialRV)
-def betabinomial_mean(op, rv, rng, size, dtype, n, alpha, beta):
+def betabinomial_mean(op, rv, rng, size, n, alpha, beta):
     return maybe_resize((n * alpha) / (alpha + beta), size)
 
 
 @_mean.register(BetaClippedRV)
-def beta_clipped_mean(op, rv, rng, size, dtype, alpha, beta):
+def beta_clipped_mean(op, rv, rng, size, alpha, beta):
     return maybe_resize(alpha / (alpha + beta), size)
 
 
 @_mean.register(BetaRV)
-def beta_mean(op, rv, rng, size, dtype, alpha, beta):
+def beta_mean(op, rv, rng, size, alpha, beta):
     return maybe_resize(alpha / (alpha + beta), size)
 
 
 @_mean.register(BinomialRV)
-def binomial_mean(op, rv, rng, size, dtype, n, p):
+def binomial_mean(op, rv, rng, size, n, p):
     return maybe_resize(n * p, size)
 
 
@@ -132,7 +132,7 @@ def binomial_mean(op, rv, rng, size, dtype, n, p):
 
 
 @_mean.register(DirichletRV)
-def dirichlet_mean(op, rv, rng, size, dtype, a):
+def dirichlet_mean(op, rv, rng, size, a):
     norm_constant = pt.sum(a, axis=-1)[..., None]
     mean = a / norm_constant
     if not rv_size_is_none(size):
@@ -141,60 +141,60 @@ def dirichlet_mean(op, rv, rng, size, dtype, a):
 
 
 @_mean.register(ExponentialRV)
-def exponential_mean(op, rv, rng, size, dtype, mu):
+def exponential_mean(op, rv, rng, size, mu):
     return maybe_resize(mu, size)
 
 
 @_mean.register(GammaRV)
-def gamma_mean(op, rv, rng, size, dtype, alpha, inv_beta):
+def gamma_mean(op, rv, rng, size, alpha, inv_beta):
     # The Aesara `GammaRV` `Op` inverts the `beta` parameter itself
     return maybe_resize(alpha * inv_beta, size)
 
 
 @_mean.register(GeometricRV)
-def geometric_mean(op, rv, rng, size, dtype, p):
+def geometric_mean(op, rv, rng, size, p):
     return maybe_resize(1.0 / p, size)
 
 
 @_mean.register(GumbelRV)
-def gumbel_mean(op, rv, rng, size, dtype, mu, beta):
+def gumbel_mean(op, rv, rng, size, mu, beta):
     return maybe_resize(mu + beta * np.euler_gamma, size)
 
 
 @_mean.register(HalfNormalRV)
-def halfnormal_mean(op, rv, rng, size, dtype, loc, sigma):
+def halfnormal_mean(op, rv, rng, size, loc, sigma):
     _, sigma = pt.broadcast_arrays(loc, sigma)
     return maybe_resize(sigma * pt.sqrt(2 / np.pi), size)
 
 
 @_mean.register(HyperGeometricRV)
-def hypergeometric_mean(op, rv, rng, size, dtype, good, bad, n):
+def hypergeometric_mean(op, rv, rng, size, good, bad, n):
     N, k = good + bad, good
     return maybe_resize(n * k / N, size)
 
 
 @_mean.register(InvGammaRV)
-def invgamma_mean(op, rv, rng, size, dtype, alpha, beta):
+def invgamma_mean(op, rv, rng, size, alpha, beta):
     return maybe_resize(pt.switch(alpha > 1, beta / (alpha - 1.0), np.nan), size)
 
 
 @_mean.register(LaplaceRV)
-def laplace_mean(op, rv, rng, size, dtype, mu, b):
+def laplace_mean(op, rv, rng, size, mu, b):
     return maybe_resize(pt.broadcast_arrays(mu, b)[0], size)
 
 
 @_mean.register(LogisticRV)
-def logistic_mean(op, rv, rng, size, dtype, mu, s):
+def logistic_mean(op, rv, rng, size, mu, s):
     return maybe_resize(pt.broadcast_arrays(mu, s)[0], size)
 
 
 @_mean.register(LogNormalRV)
-def lognormal_mean(op, rv, rng, size, dtype, mu, sigma):
+def lognormal_mean(op, rv, rng, size, mu, sigma):
     return maybe_resize(pt.exp(mu + 0.5 * sigma**2), size)
 
 
 @_mean.register(MultinomialRV)
-def multinomial_mean(op, rv, rng, size, dtype, n, p):
+def multinomial_mean(op, rv, rng, size, n, p):
     n = pt.shape_padright(n)
     mean = n * p
     if not rv_size_is_none(size):
@@ -204,7 +204,7 @@ def multinomial_mean(op, rv, rng, size, dtype, n, p):
 
 
 @_mean.register(MvNormalRV)
-def mvnormal_mean(op, rv, rng, size, dtype, mu, cov):
+def mvnormal_mean(op, rv, rng, size, mu, cov):
     mean = mu
     if not rv_size_is_none(size):
         mean_size = pt.concatenate([size, [mu.shape[-1]]])
@@ -213,70 +213,70 @@ def mvnormal_mean(op, rv, rng, size, dtype, mu, cov):
 
 
 @_mean.register(NegBinomialRV)
-def negative_binomial_mean(op, rv, rng, size, dtype, n, p):
+def negative_binomial_mean(op, rv, rng, size, n, p):
     return maybe_resize(n * (1 - p) / p, size)
 
 
 @_mean.register(NormalRV)
-def normal_mean(op, rv, rng, size, dtype, mu, sigma):
+def normal_mean(op, rv, rng, size, mu, sigma):
     return maybe_resize(pt.broadcast_arrays(mu, sigma)[0], size)
 
 
 @_mean.register(ParetoRV)
-def pareto_mean(op, rv, rng, size, dtype, alpha, m):
+def pareto_mean(op, rv, rng, size, alpha, m):
     return maybe_resize(pt.switch(alpha > 1, alpha * m / (alpha - 1), np.nan), size)
 
 
 @_mean.register(PoissonRV)
-def poisson_mean(op, rv, rng, size, dtype, mu):
+def poisson_mean(op, rv, rng, size, mu):
     return maybe_resize(mu, size)
 
 
 @_mean.register(TriangularRV)
-def triangular_mean(op, rv, rng, size, dtype, lower, c, upper):
+def triangular_mean(op, rv, rng, size, lower, c, upper):
     return maybe_resize((lower + upper + c) / 3, size)
 
 
 @_mean.register(UniformRV)
-def uniform_mean(op, rv, rng, size, dtype, lower, upper):
+def uniform_mean(op, rv, rng, size, lower, upper):
     return maybe_resize((lower + upper) / 2, size)
 
 
 @_mean.register(VonMisesRV)
-def vonmisses_mean(op, rv, rng, size, dtype, mu, kappa):
+def vonmisses_mean(op, rv, rng, size, mu, kappa):
     return maybe_resize(pt.broadcast_arrays(mu, kappa)[0], size)
 
 
 @_mean.register(KumaraswamyRV)
-def kumaraswamy_mean(op, rv, rng, size, dtype, a, b):
+def kumaraswamy_mean(op, rv, rng, size, a, b):
     return maybe_resize(
         pt.exp(pt.log(b) + pt.gammaln(1 + 1 / a) + pt.gammaln(b) - pt.gammaln(1 + 1 / a + b)),
         size,
     )
 
 
 @_mean.register(WaldRV)
-def wald_mean(op, rv, rng, size, dtype, mu, lam, alpha):
+def wald_mean(op, rv, rng, size, mu, lam, alpha):
     return maybe_resize(pt.broadcast_arrays(mu, lam, alpha)[0], size)
 
 
 @_mean.register(WeibullBetaRV)
-def weibull_mean(op, rv, rng, size, dtype, alpha, beta):
+def weibull_mean(op, rv, rng, size, alpha, beta):
     return maybe_resize(beta * pt.gamma(1 + 1 / alpha), size)
 
 
 @_mean.register(AsymmetricLaplaceRV)
-def asymmetric_laplace_mean(op, rv, rng, size, dtype, b, kappa, mu):
+def asymmetric_laplace_mean(op, rv, rng, size, b, kappa, mu):
     return maybe_resize(mu - (kappa - 1 / kappa) / b, size)
 
 
 @_mean.register(StudentTRV)
-def studentt_mean(op, rv, rng, size, dtype, nu, mu, sigma):
+def studentt_mean(op, rv, rng, size, nu, mu, sigma):
     return maybe_resize(pt.broadcast_arrays(mu, nu, sigma)[0], size)
 
 
 @_mean.register(HalfStudentTRV)
-def half_studentt_mean(op, rv, rng, size, dtype, nu, sigma):
+def half_studentt_mean(op, rv, rng, size, nu, sigma):
     return maybe_resize(
         pt.switch(
             nu > 1,
@@ -291,18 +291,18 @@ def half_studentt_mean(op, rv, rng, size, dtype, nu, sigma):
 
 
 @_mean.register(ExGaussianRV)
-def exgaussian_mean(op, rv, rng, size, dtype, mu, nu, sigma):
+def exgaussian_mean(op, rv, rng, size, mu, nu, sigma):
     mu, nu, _ = pt.broadcast_arrays(mu, nu, sigma)
     return maybe_resize(mu + nu, size)
 
 
 @_mean.register(SkewNormalRV)
-def skew_normal_mean(op, rv, rng, size, dtype, mu, sigma, alpha):
+def skew_normal_mean(op, rv, rng, size, mu, sigma, alpha):
     return maybe_resize(mu + sigma * (2 / np.pi) ** 0.5 * alpha / (1 + alpha**2) ** 0.5, size)
 
 
 @_mean.register(RiceRV)
-def rice_mean(op, rv, rng, size, dtype, nu, sigma):
+def rice_mean(op, rv, rng, size, nu, sigma):
     nu_sigma_ratio = -(nu**2) / (2 * sigma**2)
     return maybe_resize(
         sigma
@@ -317,22 +317,22 @@ def rice_mean(op, rv, rng, size, dtype, nu, sigma):
 
 
 @_mean.register(MoyalRV)
-def moyal_mean(op, rv, rng, size, dtype, mu, sigma):
+def moyal_mean(op, rv, rng, size, mu, sigma):
     return maybe_resize(mu + sigma * (np.euler_gamma + pt.log(2)), size)
 
 
 @_mean.register(PolyaGammaRV)
-def polya_gamma_mean(op, rv, rng, size, dtype, h, z):
+def polya_gamma_mean(op, rv, rng, size, h, z):
     return maybe_resize(pt.switch(pt.eq(z, 0), h / 4, tanh(z / 2) * (h / (2 * z))), size)
 
 
 @_mean.register(DiscreteUniformRV)
-def discrete_uniform_mean(op, rv, rng, size, dtype, lower, upper):
+def discrete_uniform_mean(op, rv, rng, size, lower, upper):
     return maybe_resize((upper + lower) / 2.0, size)
 
 
 @_mean.register(DiracDeltaRV)
-def dirac_delta_mean(op, rv, rng, size, dtype, c):
+def dirac_delta_mean(op, rv, rng, size, c):
     return maybe_resize(c, size)
 
 
@@ -355,12 +355,12 @@ def mixture_mean(op, rv, rng, weights, *components):
 
 
 @_mean.register(GeneralizedPoissonRV)
-def generalized_poisson_mean(op, rv, rng, size, dtype, mu, lam):
+def generalized_poisson_mean(op, rv, rng, size, mu, lam):
     return maybe_resize(mu / (1 - lam), size)
 
 
 @_mean.register(MvStudentTRV)
-def mvstudentt_mean(op, rv, rng, size, dtype, nu, mu, scale):
+def mvstudentt_mean(op, rv, rng, size, nu, mu, scale):
     mean = mu
     if not rv_size_is_none(size):
         mean_size = pt.concatenate([size, [mu.shape[-1]]])
@@ -369,7 +369,7 @@ def mvstudentt_mean(op, rv, rng, size, dtype, nu, mu, scale):
 
 
 @_mean.register(DirichletMultinomialRV)
-def dirichlet_multinomial_mean(op, rv, rng, size, dtype, n, a):
+def dirichlet_multinomial_mean(op, rv, rng, size, n, a):
     mean = pt.shape_padright(n) * a / pt.sum(a, axis=-1, keepdims=True)
     if not rv_size_is_none(size):
         output_size = pt.concatenate([size, [a.shape[-1]]])
@@ -386,17 +386,17 @@ def lkj_cholesky_cov_mean(op, rv, rng, n, eta, sd_dist):
 
 
 @_mean.register(LKJCorrRV)
-def lkj_corr_mean(op, rv, rng, size, dtype, *args):
+def lkj_corr_mean(op, rv, rng, size, *args):
     return pt.full_like(rv, pt.eye(rv.shape[-1]))
 
 
 @_mean.register(MatrixNormalRV)
-def matrix_normal_mean(op, rv, rng, size, dtype, mu, rowchol, colchol):
+def matrix_normal_mean(op, rv, rng, size, mu, rowchol, colchol):
     return pt.full_like(rv, mu)
 
 
 @_mean.register(KroneckerNormalRV)
-def kronecker_normal_mean(op, rv, rng, size, dtype, mu, covs, chols, evds):
+def kronecker_normal_mean(op, rv, rng, size, mu, covs, chols, evds):
     mean = mu
     if not rv_size_is_none(size):
         mean_size = pt.concatenate([size, mu.shape])
@@ -405,12 +405,12 @@ def kronecker_normal_mean(op, rv, rng, size, dtype, mu, covs, chols, evds):
 
 
 @_mean.register(CARRV)
-def car_mean(op, rv, rng, size, dtype, mu, W, alpha, tau):
+def car_mean(op, rv, rng, size, mu, W, alpha, tau):
     return pt.full_like(rv, mu)
 
 
 @_mean.register(StickBreakingWeightsRV)
-def stick_breaking_mean(op, rv, rng, size, dtype, alpha, K):
+def stick_breaking_mean(op, rv, rng, size, alpha, K):
     alpha = alpha[..., np.newaxis]
     mean = (alpha / (1 + alpha)) ** pt.arange(K)
     mean *= 1 / (1 + alpha)