Тензор потока: Как найти диагональный элемент тензора ранга 4 в тензор ранга 3

Question

Предположим, у меня есть тензор четвертого ранга: T_{i,j,p,q},, и я хочу найти диагональный элемент первых двух измерений, чтобы сформировать новый тензор: T_{i,i,p,q} в T_{i,p,q}.Как мне это сделать в tenorflow?

Спасибо