warp.tile_broadcast

warp.tile_broadcast(

a: Tile[Any, tuple[int, ...]],

shape: tuple[int, ...],

) → Tile[Any, tuple[int, ...]]

• Kernel: true
• Python: false
• Differentiable: true

Broadcast a tile to a larger shape.

Broadcasting follows numpy.broadcast_to(): the shapes are aligned from the right, each source dimension must either match the target or have length one, and leading dimensions may be added. It is one-way, to the explicit target shape, which must have between one and four dimensions.

The result aliases a instead of copying it. The current implementation places a in shared memory; the result is non-owning and allocates no additional storage. The result is writable, but every position along a broadcast dimension refers to the same element of a, so a write updates all of them. Concurrent or collective writes of different values through aliased positions race; treat the view as read-only unless each underlying element has exactly one writer.

Parameters:

• a – Tile to broadcast

• shape – The shape to broadcast to, whose entries must be compile-time constants and which must have at least as many dimensions as a

Returns:

A non-owning tile with the broadcast shape that aliases a.

Example

@wp.kernel
def repeat_row(a: wp.array[float], out: wp.array2d[float]):
    t = wp.tile_load(a, shape=3)
    b = wp.tile_broadcast(t, shape=(2, 3))

    wp.tile_store(out, b)

a = wp.array([1.0, 2.0, 3.0], dtype=float)
out = wp.zeros((2, 3), dtype=float)

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

print(out.numpy())

[[1. 2. 3.]
 [1. 2. 3.]]