pymc.math.stack

pymc.math.stack(tensors, axis=0)

Stack tensors in sequence on given axis (default is 0).

Take a sequence of tensors or tensor-like constant and stack them on given axis to make a single tensor. The size in dimension axis of the result will be equal to the number of tensors passed.

Parameters:

tensorsSequence[TensorLike]

A list of tensors or tensor-like constants to be stacked.

axisint

The index of the new axis. Default value is 0.

Examples

>>> a = pytensor.tensor.type.scalar()
>>> b = pytensor.tensor.type.scalar()
>>> c = pytensor.tensor.type.scalar()
>>> x = pytensor.tensor.stack([a, b, c])
>>> x.ndim  # x is a vector of length 3.
1
>>> a = pytensor.tensor.type.tensor4()
>>> b = pytensor.tensor.type.tensor4()
>>> c = pytensor.tensor.type.tensor4()
>>> x = pytensor.tensor.stack([a, b, c])
>>> x.ndim  # x is a 5d tensor.
5
>>> rval = x.eval(dict((t, np.zeros((2, 2, 2, 2))) for t in [a, b, c]))
>>> rval.shape  # 3 tensors are stacked on axis 0
(3, 2, 2, 2, 2)
>>> x = pytensor.tensor.stack([a, b, c], axis=3)
>>> x.ndim
5
>>> rval = x.eval(dict((t, np.zeros((2, 2, 2, 2))) for t in [a, b, c]))
>>> rval.shape  # 3 tensors are stacked on axis 3
(2, 2, 2, 3, 2)
>>> x = pytensor.tensor.stack([a, b, c], axis=-2)
>>> x.ndim
5
>>> rval = x.eval(dict((t, np.zeros((2, 2, 2, 2))) for t in [a, b, c]))
>>> rval.shape  # 3 tensors are stacked on axis -2
(2, 2, 2, 3, 2)