Oh. Nice. I tried multiplying ‘a’ by one hot vectors:
index_tensor = index.view(*index.size(), -1)
one_hot = torch.zeros(*index.size(), 2, dtype=index.dtype)
one_hot = one_hot.scatter(1, index_tensor, 1).unsqueeze(-1)
result = a.matmul(one_hot).squeeze(-1).transpose(-2, -1)

Thinking too complicated always makes trouble!!
Thank you