ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more


Returns the diagonal part of the tensor.

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_part produce 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]],
     [[[2111, 2112], [2121, 2122]],
      [[2211, 2212], [2221, 2222]]]
<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
array([[1111, 1212],
       [2121, 2222]], dtype=int32)>
TensorShape([2, 2, 2])

input A Tensor with rank 2k.
name A name for the operation (optional).

A Tensor containing diagonals of input. Has the same type as input, and rank k.