shape_utils

This submodule contains various functions that apply numpy’s broadcasting rules to shape tuples, and also to samples drawn from probability distributions.

The main challenge when broadcasting samples drawn from a generative model, is that each random variate has a core shape. When we draw many i.i.d samples from a given RV, for example if we ask for size_tuple i.i.d draws, the result usually is a size_tuple + RV_core_shape. In the generative model’s hierarchy, the downstream RVs that are conditionally dependent on our above sampled values, will get an array with a shape that is incosistent with the core shape they expect to see for their parameters. This is a problem sometimes because it prevents regular broadcasting in complex hierachical models, and thus make prior and posterior predictive sampling difficult.

This module introduces functions that are made aware of the requested size_tuple of i.i.d samples, and does the broadcasting on the core shapes, transparently ignoring or moving the i.i.d size_tuple prepended axes around.

to_tuple(shape)

Convert ints, arrays, and Nones to tuples

shapes_broadcasting(*args[, raise_exception])

Return the shape resulting from broadcasting multiple shapes.

broadcast_dist_samples_shape(shapes[, size])

Apply shape broadcasting to shape tuples but assuming that the shapes correspond to draws from random variables, with the size tuple possibly prepended to it.

get_broadcastable_dist_samples(samples[, …])

Get a view of the samples drawn from distributions which adds new axises in between the size prepend and the distribution’s shape.

broadcast_distribution_samples(samples[, size])

Broadcast samples drawn from distributions taking into account the size (i.e.

broadcast_dist_samples_to(to_shape, samples)

Broadcast samples drawn from distributions to a given shape, taking into account the size (i.e.

A collection of common numpy array shape operations needed for broadcasting samples from probability distributions for stochastic nodes in PyMC.

pymc3.distributions.shape_utils.broadcast_dist_samples_shape(shapes, size=None)

Apply shape broadcasting to shape tuples but assuming that the shapes correspond to draws from random variables, with the size tuple possibly prepended to it. The size prepend is ignored to consider if the supplied shapes can broadcast or not. It is prepended to the resulting broadcasted shapes, if any of the shape tuples had the size prepend.

Parameters
shapes: Iterable of tuples holding the distribution samples shapes
size: None, int or tuple (optional)

size of the sample set requested.

Returns
tuple of the resulting shape

Examples

size = 100
shape0 = (size,)
shape1 = (size, 5)
shape2 = (size, 4, 5)
out = broadcast_dist_samples_shape([shape0, shape1, shape2],
                                   size=size)
assert out == (size, 4, 5)
size = 100
shape0 = (size,)
shape1 = (5,)
shape2 = (4, 5)
out = broadcast_dist_samples_shape([shape0, shape1, shape2],
                                   size=size)
assert out == (size, 4, 5)
size = 100
shape0 = (1,)
shape1 = (5,)
shape2 = (4, 5)
out = broadcast_dist_samples_shape([shape0, shape1, shape2],
                                   size=size)
assert out == (4, 5)
pymc3.distributions.shape_utils.broadcast_dist_samples_to(to_shape, samples, size=None)

Broadcast samples drawn from distributions to a given shape, taking into account the size (i.e. the number of samples) of the draw, which is prepended to the sample’s shape.

Parameters
to_shape: Tuple shape onto which the samples must be able to broadcast
samples: Iterable of ndarrays holding the sampled values
size: None, int or tuple (optional)

size of the sample set requested.

Returns
List of the broadcasted sample arrays

Examples

to_shape = (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(
    to_shape,
    [sample0, sample1, sample2],
    size=size
)
assert np.all((o.shape == (size, 3, 4, 5) for o in out))
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
to_shape = (3, 1, 5)
sample0 = np.random.randn(size)
sample1 = np.random.randn(5)
sample2 = np.random.randn(4, 5)
out = broadcast_dist_samples_to(
    to_shape,
    [sample0, sample1, sample2],
    size=size
)
assert np.all((o.shape == (size, 3, 4, 5) for o in out))
assert np.all(sample0[:, None, None, None] == out[0])
assert np.all(sample1 == out[1])
assert np.all(sample2 == out[2])
pymc3.distributions.shape_utils.broadcast_distribution_samples(samples, size=None)

Broadcast samples drawn from distributions taking into account the size (i.e. the number of samples) of the draw, which is prepended to the sample’s shape.

Parameters
samples: Iterable of ndarrays holding the sampled values
size: None, int or tuple (optional)

size of the sample set requested.

Returns
List of broadcasted sample arrays

Examples

size = 100
sample0 = np.random.randn(size)
sample1 = np.random.randn(size, 5)
sample2 = np.random.randn(size, 4, 5)
out = broadcast_distribution_samples([sample0, sample1, sample2],
                                     size=size)
assert all((o.shape == (size, 4, 5) for o in out))
assert np.all(sample0[:, None, None] == out[0])
assert np.all(sample1[:, None, :] == out[1])
assert np.all(sample2 == out[2])
size = 100
sample0 = np.random.randn(size)
sample1 = np.random.randn(5)
sample2 = np.random.randn(4, 5)
out = broadcast_distribution_samples([sample0, sample1, sample2],
                                     size=size)
assert all((o.shape == (size, 4, 5) for o in out))
assert np.all(sample0[:, None, None] == out[0])
assert np.all(sample1 == out[1])
assert np.all(sample2 == out[2])
pymc3.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 axises 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])
pymc3.distributions.shape_utils.shapes_broadcasting(*args, raise_exception=False)

Return the shape resulting from broadcasting multiple shapes. Represents numpy’s broadcasting rules.

Parameters
*args: array-like of int

Tuples or arrays or lists representing the shapes of arrays to be broadcast.

raise_exception: bool (optional)

Controls whether to raise an exception or simply return None if the broadcasting fails.

Returns
Resulting shape. If broadcasting is not possible and raise_exception is
False, then None is returned. If raise_exception is True, a
ValueError is raised.
pymc3.distributions.shape_utils.to_tuple(shape)

Convert ints, arrays, and Nones to tuples

Parameters
shape: None, int or array-like

Represents the shape to convert to tuple.

Returns
If shape is None, returns an empty tuple. If it’s an int, (shape,) is
returned. If it is array-like, tuple(shape) is returned.