warp.tile_squeeze#

warp.tile_squeeze( t: Tile[Any, tuple[int, ...]], axis: tuple[int, ...] = ..., ) → Tile[Any, tuple[int, ...]]#

Kernel: true
Python: false
Differentiable: true

Return a view of a tile with dimensions of length one removed.

The view aliases the source tile, so writing through it modifies t. Taking the view requires shared storage, but the view itself is non-owning and allocates no additional storage.

Parameters:

t – Input tile to squeeze
axis – Axis or axes to remove, as a compile-time constant. If omitted, all dimensions of t with extent one are removed. Each specified axis must refer to a dimension of t with extent one. Negative axes count from the end.

Returns:

A non-owning tile that aliases t with all, or the selected, dimensions of length one removed.

Example

@wp.kernel
def remove_singleton_dims(a: wp.array3d[float], out: wp.array[float]):
    t = wp.tile_load(a, shape=(1, 4, 1))
    s = wp.tile_squeeze(t)

    wp.tile_store(out, s)

a = wp.array(np.arange(1, 5, dtype=np.float32).reshape(1, 4, 1), dtype=float)
out = wp.zeros(4, dtype=float)

wp.launch_tiled(remove_singleton_dims, dim=1, inputs=[a], outputs=[out], block_dim=2)

print(out.numpy())

[1. 2. 3. 4.]