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

tf.linalg.tensor_diag_part

Stay organized with collections Save and categorize content based on your preferences.

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])