Given a tensor like
a =  1  1  1
     1  1  1
     1  1  1
     2  1  1
     2  1  1

I would like to obtain all rows starting with 2, i.e. I would like to obtain
b = 2 1 1 
    2 1 1

Indexing using a[a[:, 0]==2] doesn’t seem to work and index_select requires indices to be known and so does gather.