Skip to content

Commit

Permalink
Get rid of uses of intX besides convert_observed_data
Browse files Browse the repository at this point in the history
  • Loading branch information
ricardoV94 committed May 23, 2024
1 parent eddf03b commit 85a1e36
Show file tree
Hide file tree
Showing 5 changed files with 15 additions and 16 deletions.
6 changes: 3 additions & 3 deletions pymc/distributions/multivariate.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@
from pymc.distributions.transforms import Interval, ZeroSumTransform, _default_transform
from pymc.logprob.abstract import _logprob
from pymc.math import kron_diag, kron_dot
from pymc.pytensorf import intX, normalize_rng_param
from pymc.pytensorf import normalize_rng_param
from pymc.util import check_dist_not_registered

__all__ = [
Expand Down Expand Up @@ -929,7 +929,7 @@ class Wishart(Continuous):

@classmethod
def dist(cls, nu, V, *args, **kwargs):
nu = pt.as_tensor_variable(intX(nu))
nu = pt.as_tensor_variable(nu, dtype=int)
V = pt.as_tensor_variable(V)

warnings.warn(
Expand Down Expand Up @@ -2454,7 +2454,7 @@ class StickBreakingWeightsRV(RandomVariable):

def make_node(self, rng, size, dtype, alpha, K):
alpha = pt.as_tensor_variable(alpha)
K = pt.as_tensor_variable(intX(K))
K = pt.as_tensor_variable(K, dtype=int)

if K.ndim > 0:
raise ValueError("K must be a scalar.")
Expand Down
8 changes: 4 additions & 4 deletions pymc/distributions/timeseries.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@
from pymc.exceptions import NotConstantValueError
from pymc.logprob.abstract import _logprob
from pymc.logprob.basic import logp
from pymc.pytensorf import constant_fold, intX
from pymc.pytensorf import constant_fold
from pymc.util import check_dist_not_registered

__all__ = [
Expand Down Expand Up @@ -174,7 +174,7 @@ def dist(cls, init_dist, innovation_dist, steps=None, **kwargs) -> pt.TensorVari
)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = pt.as_tensor_variable(intX(steps))
steps = pt.as_tensor_variable(steps, dtype=int)

return super().dist([init_dist, innovation_dist, steps], **kwargs)

Expand Down Expand Up @@ -599,7 +599,7 @@ def dist(
)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = pt.as_tensor_variable(intX(steps), ndim=0)
steps = pt.as_tensor_variable(steps, dtype=int, ndim=0)

if init_dist is not None:
if not isinstance(init_dist, TensorVariable) or not isinstance(
Expand Down Expand Up @@ -961,7 +961,7 @@ def dist(cls, dt, sde_fn, sde_pars, *, init_dist=None, steps=None, **kwargs):
)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = pt.as_tensor_variable(intX(steps), ndim=0)
steps = pt.as_tensor_variable(steps, dtype=int, ndim=0)

dt = pt.as_tensor_variable(dt)
sde_pars = [pt.as_tensor_variable(x) for x in sde_pars]
Expand Down
4 changes: 2 additions & 2 deletions pymc/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
local_check_parameter_to_ninf_switch,
rvs_in_graph,
)
from pymc.pytensorf import compile_pymc, floatX, inputvars, intX
from pymc.pytensorf import compile_pymc, floatX, inputvars

# This mode can be used for tests where model compilations takes the bulk of the runtime
# AND where we don't care about posterior numerical or sampling stability (e.g., when
Expand Down Expand Up @@ -771,7 +771,7 @@ def discrete_random_tester(
f = fails
while p <= alpha and f > 0:
o = pymc_rand()
e = intX(ref_rand(size=size, **point))
e = ref_rand(size=size, **point).astype(int)
o = np.atleast_1d(o).flatten()
e = np.atleast_1d(e).flatten()
bins = min(20, max(len(set(e)), len(set(o))))
Expand Down
10 changes: 5 additions & 5 deletions tests/distributions/test_multivariate.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
from pymc.logprob.basic import logp
from pymc.logprob.utils import ParameterValueError
from pymc.math import kronecker
from pymc.pytensorf import compile_pymc, floatX, intX
from pymc.pytensorf import compile_pymc, floatX
from pymc.sampling.forward import draw
from pymc.testing import (
BaseTestDistributionRandom,
Expand Down Expand Up @@ -674,8 +674,8 @@ def test_multinomial_p_not_normalized_symbolic(self):
)
@pytest.mark.parametrize("extra_size", [(1,), (2,), (2, 3)])
def test_multinomial_vectorized(self, n, p, extra_size):
n = intX(np.array(n))
p = floatX(np.array(p))
n = np.array(n)
p = np.array(p)
p /= p.sum(axis=-1, keepdims=True)

_, bcast_p = broadcast_params([n, p], ndims_params=[0, 1])
Expand Down Expand Up @@ -757,8 +757,8 @@ def test_dirichlet_multinomial_matches_beta_binomial(self):
)
@pytest.mark.parametrize("extra_size", [(1,), (2,), (2, 3)])
def test_dirichlet_multinomial_vectorized(self, n, a, extra_size):
n = intX(np.array(n))
a = floatX(np.array(a))
n = np.array(n)
a = np.array(a)

_, bcast_a = broadcast_params([n, a], ndims_params=[0, 1])
size = extra_size + bcast_a.shape[:-1]
Expand Down
3 changes: 1 addition & 2 deletions tests/variational/test_inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
import pymc as pm
import pymc.variational.opvi as opvi

from pymc.pytensorf import intX
from pymc.variational.inference import ADVI, ASVGD, SVGD, FullRankADVI
from pymc.variational.opvi import NotImplementedInference
from tests import models
Expand Down Expand Up @@ -278,7 +277,7 @@ def test_profile(inference):
@pytest.fixture(scope="module")
def binomial_model():
n_samples = 100
xs = intX(np.random.binomial(n=1, p=0.2, size=n_samples))
xs = np.random.binomial(n=1, p=0.2, size=n_samples)
with pm.Model() as model:
p = pm.Beta("p", alpha=1, beta=1)
pm.Binomial("xs", n=1, p=p, observed=xs)
Expand Down

0 comments on commit 85a1e36

Please sign in to comment.