This document is relevant for: Inf2, Trn1, Trn1n

nki.language.expand_dims#

nki.language.expand_dims(data, axis)[source]#

Expand the shape of a tile. Insert a new axis that will appear at the axis position in the expanded tile shape. Currently only supports expanding dimensions after the last index of the tile.

((Similar to numpy.expand_dims))

Parameters:

  • data – a tile input

  • axis – int or tuple/list of ints. Position in the expanded axes where the new axis (or axes) is placed; must be free dimensions, not partition dimension (0); Currently only supports axis (or axes) after the last index.

Returns:

a tile with view of input data with the number of dimensions increased.

This document is relevant for: Inf2, Trn1, Trn1n