tf.eye

View source on GitHub

Construct an identity matrix, or a batch of matrices.

tf.eye(
    num_rows, num_columns=None, batch_shape=None, dtype=tf.dtypes.float32, name=None
)
# Construct one identity matrix.
tf.eye(2)
==> [[1., 0.],
     [0., 1.]]

# Construct a batch of 3 identity matricies, each 2 x 2.
# batch_identity[i, :, :] is a 2 x 2 identity matrix, i = 0, 1, 2.
batch_identity = tf.eye(2, batch_shape=[3])

# Construct one 2 x 3 "identity" matrix
tf.eye(2, num_columns=3)
==> [[ 1.,  0.,  0.],
     [ 0.,  1.,  0.]]

Args:

Returns:

A Tensor of shape batch_shape + [num_rows, num_columns]