tf.keras.ops.diag | TensorFlow v2.16.1 (original) (raw)
Extract a diagonal or construct a diagonal array.
View aliases
Main aliases
tf.keras.ops.diag(
x, k=0
)
| Args | |
|---|---|
| x | Input tensor. If x is 2-D, returns the k-th diagonal of x. If x is 1-D, return a 2-D tensor with x on the k-th diagonal. |
| k | The diagonal to consider. Defaults to 0. Use k > 0 for diagonals above the main diagonal, and k < 0 for diagonals below the main diagonal. |
| Returns |
|---|
| The extracted diagonal or constructed diagonal tensor. |
Examples:
from keras.src import ops
x = ops.arange(9).reshape((3, 3))
x
array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
ops.diag(x)
array([0, 4, 8])
ops.diag(x, k=1)
array([1, 5])
ops.diag(x, k=-1)
array([3, 7])
ops.diag(ops.diag(x)))
array([[0, 0, 0],
[0, 4, 0],
[0, 0, 8]])