Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix size issue when using broadcastable dimensions in multivariate RandomVariables #1056

Merged
merged 1 commit into from
Jul 17, 2022
Merged

Fix size issue when using broadcastable dimensions in multivariate RandomVariables #1056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 8 additions & 27 deletions aesara/tensor/random/basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,19 +331,13 @@ def rng_fn(cls, rng, mean, cov, size):
# Neither SciPy nor NumPy implement parameter broadcasting for
# multivariate normals (or any other multivariate distributions),
# so we need to implement that here
mean, cov = broadcast_params([mean, cov], cls.ndims_params)
size = tuple(size or ())

size = tuple(size or ())
if size:
if (
0 < mean.ndim - 1 <= len(size)
and size[-mean.ndim + 1 :] != mean.shape[:-1]
):
raise ValueError(
"shape mismatch: objects cannot be broadcast to a single shape"
)
mean = np.broadcast_to(mean, size + mean.shape[-1:])
cov = np.broadcast_to(cov, size + cov.shape[-2:])
else:
mean, cov = broadcast_params([mean, cov], cls.ndims_params)

res = np.empty(mean.shape)
for idx in np.ndindex(mean.shape[:-1]):
Expand Down Expand Up @@ -374,22 +368,12 @@ def rng_fn(cls, rng, alphas, size):
size = tuple(np.atleast_1d(size))

if size:
if (
0 < alphas.ndim - 1 <= len(size)
and size[-alphas.ndim + 1 :] != alphas.shape[:-1]
):
raise ValueError(
"shape mismatch: objects cannot be broadcast to a single shape"
)
samples_shape = size + alphas.shape[-1:]
else:
samples_shape = alphas.shape
alphas = np.broadcast_to(alphas, size + alphas.shape[-1:])

samples_shape = alphas.shape
samples = np.empty(samples_shape)
alphas_bcast = np.broadcast_to(alphas, samples_shape)

for index in np.ndindex(*samples_shape[:-1]):
samples[index] = rng.dirichlet(alphas_bcast[index])
samples[index] = rng.dirichlet(alphas[index])

return samples
else:
Expand Down Expand Up @@ -608,16 +592,13 @@ def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=Non
@classmethod
def rng_fn(cls, rng, n, p, size):
if n.ndim > 0 or p.ndim > 1:
n, p = broadcast_params([n, p], cls.ndims_params)
size = tuple(size or ())

if size:
if 0 < p.ndim - 1 <= len(size) and size[-p.ndim + 1 :] != p.shape[:-1]:
raise ValueError(
"shape mismatch: objects cannot be broadcast to a single shape"
)
n = np.broadcast_to(n, size)
p = np.broadcast_to(p, size + p.shape[-1:])
else:
n, p = broadcast_params([n, p], cls.ndims_params)

res = np.empty(p.shape, dtype=cls.dtype)
for idx in np.ndindex(p.shape[:-1]):
Expand Down
28 changes: 20 additions & 8 deletions tests/tensor/random/test_basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -578,9 +578,9 @@ def test_mvnormal_samples(mu, cov, size):
def test_mvnormal_default_args():
compare_sample_values(multivariate_normal, test_fn=mvnormal_test_fn)

with pytest.raises(ValueError, match="shape mismatch.*"):
with pytest.raises(ValueError, match="operands could not be broadcast together "):
multivariate_normal.rng_fn(
None, np.zeros((1, 2)), np.ones((1, 2, 2)), size=(4,)
None, np.zeros((3, 2)), np.ones((3, 2, 2)), size=(4,)
)


Expand Down Expand Up @@ -654,11 +654,17 @@ def dirichlet_test_fn(mean=None, cov=None, size=None, random_state=None):
def test_dirichlet_rng():
alphas = np.array([[100, 1, 1], [1, 100, 1], [1, 1, 100]], dtype=config.floatX)

with pytest.raises(ValueError, match="shape mismatch.*"):
# The independent dimension's shape is missing from size (i.e. should
# be `(10, 2, 3)`)
with pytest.raises(ValueError, match="operands could not be broadcast together"):
# The independent dimension's shape cannot be broadcasted from (3,) to (10, 2)
dirichlet.rng_fn(None, alphas, size=(10, 2))

with pytest.raises(
ValueError, match="input operand has more dimensions than allowed"
):
# One of the independent dimension's shape is missing from size
# (i.e. should be `(1, 3)`)
dirichlet.rng_fn(None, np.broadcast_to(alphas, (1, 3, 3)), size=(3,))


M_at = iscalar("M")
M_at.tag.test_value = 3
Expand Down Expand Up @@ -1146,11 +1152,17 @@ def test_multinomial_rng():
test_M = np.array([10, 20], dtype=np.int64)
test_p = np.array([[0.999, 0.001], [0.001, 0.999]], dtype=config.floatX)

with pytest.raises(ValueError, match="shape mismatch.*"):
# The independent dimension's shape is missing from size (i.e. should
# be `(1, 2)`)
with pytest.raises(ValueError, match="operands could not be broadcast together"):
# The independent dimension's shape cannot be broadcasted from (2,) to (1,)
multinomial.rng_fn(None, test_M, test_p, size=(1,))

with pytest.raises(
ValueError, match="input operand has more dimensions than allowed"
):
# One of the independent dimension's shape is missing from size
# (i.e. should be `(5, 2)`)
multinomial.rng_fn(None, np.broadcast_to(test_M, (5, 2)), test_p, size=(2,))


@pytest.mark.parametrize(
"p, size, test_fn",
Expand Down