mygrad.expand_dims

mygrad.expand_dims(a: ArrayLike, axis: int, *, constant: bool | None = None) → Tensor

Expand the dimensions of a tensor by adding a new axis.

This docstring was adapted from numpy.expand_dims.

Parameters:

aArrayLike

The tensor to be expanded

axisint

The position of the new axis in the expanded array shape.

constantbool, optional(default=False)

If True, the returned tensor is a constant (it does not back-propagate a gradient)

Returns:

mygrad.Tensor

Examples

>>> import mygrad as mg
>>> x = mg.Tensor([1, 2])
>>> x.shape
(2,)
>>> y = mg.expand_dims(x, 1)
>>> y.shape
(2, 1)
>>> z = mg.expand_dims(y, 0)
>>> z.shape
(1, 2, 1)