asc2.expand_dims

asc2.expand_dims(input: LocalTensor, *axis: int) → LocalTensor

Insert new dimensions of size 1 at the specified positions.

The supported data types are: int8, int16, int32, int64, float16, bfloat16, float32, float64.

Parameters:

• input – The input tensor

• axis – The positions where new dimensions should be inserted (0-based)

Returns:

A tensor with the new dimensions inserted

Return type:

LocalTensor

Raises:

TypeError – If input is not a LocalTensor

Note

Multiple axes can be specified. Axes are processed in sorted order.

Examples

Insert a dimension at axis 0:

input = asc2.copy_in(x, [0], [256])
result = input.expand_dims(0)  # shape becomes [1, 256]

Insert multiple dimensions:

input = asc2.copy_in(x, [0, 0], [32, 16])
result = input.expand_dims(0, 2)  # shape becomes [1, 32, 1, 16]

This function can also be called as a member function on LocalTensor, as input.expand_dims(...) instead of expand_dims(input, ...).