TensorFlow如何返回對(duì)角線

tf.diag_part

diag_part ( 
    input , 
    name = None
 )

參考指南：數(shù)學(xué)函數(shù)>矩陣數(shù)學(xué)函數(shù)

返回張量的對(duì)角線部分.
此操作返回輸入的對(duì)角線部分的張量.對(duì)角部分計(jì)算方法如下:
假設(shè)輸入(input)具有維度：[D1,..., dk, D1,..., dk],那么輸出是一個(gè)秩為 k 的張量,并且其維度為：[D1,..., dk] :

diagonal[i1,..., ik] = input[i1, ..., ik, i1,..., ik]

例如：

# 'input' is [[1, 0, 0, 0]
              [0, 2, 0, 0]
              [0, 0, 3, 0]
              [0, 0, 0, 4]]

tf.diag_part(input) ==> [1, 2, 3, 4]

ARGS：

• input：一個(gè)張量.必須是下列類(lèi)型之一：float32,float64,int32,int64,complex64,complex128.秩為 k 的張量的取值是 2, 4, 或6.
• name：操作的名稱(chēng)(可選).

返回：

返回張量.與 input 具有相同的類(lèi)型,為提取的對(duì)角線.