Is the combination of tensor_diag_part and gather inefficient?

Question

I have tensors A_{ij} and B_i, and I want to compute the tensor A_{iB_i}. In numpy I can do this via

A[np.arange(N), B]