tf.linalg.tensor_diag_part  |  TensorFlow v2.16.1 (original) (raw)

Returns the diagonal part of the tensor.

View aliases

Compat aliases for migration

SeeMigration guide for more details.

tf.compat.v1.diag_part

tf.linalg.tensor_diag_part(
    input, name=None
)

This operation returns a tensor with the diagonal part of the input. The diagonal part is computed as follows:

Assume input has dimensions [D1,..., Dk, D1,..., Dk], then the output is a tensor of rank k with dimensions [D1,..., Dk] where:

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

For a rank 2 tensor, linalg.diag_part and linalg.tensor_diag_partproduce the same result. For rank 3 and higher, linalg.diag_part extracts the diagonal of each inner-most matrix in the tensor. An example where they differ is given below.

x = [[[[1111,1112],[1121,1122]], [[1211,1212],[1221,1222]]], [[[2111, 2112], [2121, 2122]], [[2211, 2212], [2221, 2222]]] ] tf.linalg.tensor_diag_part(x) <tf.Tensor: shape=(2, 2), dtype=int32, numpy= array([[1111, 1212], [2121, 2222]], dtype=int32)> tf.linalg.diag_part(x).shape TensorShape([2, 2, 2])