asc2.squeeze

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

Remove dimensions of size 1 from the tensor.

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

Parameters:

• input – The input tensor.

• axis – The positions of dimensions to remove (0-based). If not provided, all dimensions of size 1 are removed.

Returns:

A tensor with the specified dimensions removed

Return type:

LocalTensor

Raises:

• TypeError – If input is not a LocalTensor

• RuntimeError – If attempting to squeeze a dimension that is not of size 1

Examples

Remove all dimensions of size 1:

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

Remove a specific dimension:

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

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