def index_select_multidim_nojit(input, indices):
    """
    Like torch.index_select, but on multiple dimensions,
    on first dimensions only (0, 1, ...)
    NB: Clamp indices if out of range

    @param input tensor of size (n_1, .., n_k, ...u)
    @param indices tensor of integers of size (...v, k)
    @return tensor of size (...v, ...u)
    """
    n = indices.size()[-1]
    input_size = input.size()

    tpl = tuple(
        indices[...,i].clamp(0, s-1)
        for i, s in enumerate(input_size[:n])
    )

    return input[(*tpl, ...)]