Best practice to index one tensor of look up tables by another tensor

I have one tensor of look up tables, say lut with shape [15, 6], where the values are all in the range of [0, 5], and can map each index to a value in each table (say the first table lut[0] will map 3 to lut[0][3]). I have one tensor of shape [7, 2, 3, 14, 14, 5] (with values in the range of [0, 5]), and I want each one of the third and last dims (corresponding to 3 and 5) to use one table for mapping. I wonder if there is some elegant and efficient way to implement this, rather than permute and list(map(*)). Thanks.

I noticed this topic, but it does not seem to be the same problem.

torch.index_select is faster than what I implemented.