helion.language.subscript

helion.language.subscript(tensor, index)[source]

Equivalent to tensor[index] where tensor is a kernel-tensor (not a host-tensor).

Can be used to add dimensions to the tensor, e.g. tensor[None, :] or tensor[:, None].

Parameters:

tensor (Tensor) – The kernel tensor to index
index (list[object]) – List of indices, including None for new dimensions and : for existing dimensions

Returns:

The indexed tensor with potentially modified dimensions

Return type:

torch.Tensor

Examples

@helion.kernel
def broadcast_multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    # x has shape (N,), y has shape (M,)
    result = torch.empty(
        [x.size(0), y.size(0)], dtype=x.dtype, device=x.device
    )

    for tile_i, tile_j in hl.tile([x.size(0), y.size(0)]):
        # Get tile data
        x_tile = x[tile_i]
        y_tile = y[tile_j]

        # Make x broadcastable: (tile_size, 1)
        # same as hl.subscript(x_tile, [slice(None), None])
        x_expanded = x_tile[:, None]
        # Make y broadcastable: (1, tile_size)
        # same as hl.subscript(y_tile, [None, slice(None)])
        y_expanded = y_tile[None, :]

        result[tile_i, tile_j] = x_expanded * y_expanded

    return result