pymc.distributions.shape_utils.get_broadcastable_dist_samples

pymc.distributions.shape_utils.get_broadcastable_dist_samples(samples, size=None, must_bcast_with=None, return_out_shape=False)

Get a view of the samples drawn from distributions which adds new axes in between the size prepend and the distribution’s shape. These views should be able to broadcast the samples from the distrubtions taking into account the size (i.e. the number of samples) of the draw, which is prepended to the sample’s shape. Optionally, one can supply an extra must_bcast_with to try to force samples to be able to broadcast with a given shape. A ValueError is raised if it is not possible to broadcast the provided samples.

Parameters

samples: Iterable of ndarrays holding the sampled values

size: None, int or tuple (optional)

size of the sample set requested.

must_bcast_with: None, int or tuple (optional)

Tuple shape to which the samples must be able to broadcast

return_out_shape: bool (optional)

If True, this function also returns the output’s shape and not only samples views.

Returns

broadcastable_samples: List of the broadcasted sample arrays

broadcast_shape: If return_out_shape is True, the resulting broadcast

shape is returned.

Examples

must_bcast_with = (3, 1, 5)
size = 100
sample0 = np.random.randn(size)
sample1 = np.random.randn(size, 5)
sample2 = np.random.randn(size, 4, 5)
out = broadcast_dist_samples_to(
    [sample0, sample1, sample2],
    size=size,
    must_bcast_with=must_bcast_with,
)
assert out[0].shape == (size, 1, 1, 1)
assert out[1].shape == (size, 1, 1, 5)
assert out[2].shape == (size, 1, 4, 5)
assert np.all(sample0[:, None, None, None] == out[0])
assert np.all(sample1[:, None, None] == out[1])
assert np.all(sample2[:, None] == out[2])

size = 100
must_bcast_with = (3, 1, 5)
sample0 = np.random.randn(size)
sample1 = np.random.randn(5)
sample2 = np.random.randn(4, 5)
out = broadcast_dist_samples_to(
    [sample0, sample1, sample2],
    size=size,
    must_bcast_with=must_bcast_with,
)
assert out[0].shape == (size, 1, 1, 1)
assert out[1].shape == (5,)
assert out[2].shape == (4, 5)
assert np.all(sample0[:, None, None, None] == out[0])
assert np.all(sample1 == out[1])
assert np.all(sample2 == out[2])