torch.eye(): Creating an Identity Matrix in PyTorch

The torch.eye() method creates a 2D tensor (identity matrix) with ones on the diagonal and zeros elsewhere. Its primary purpose is to generate identity matrices for linear algebra operations.

The above figure is a 3 x 3 identity matrix.

Syntax

torch.eye(n, 
          m=None,
          out=None, 
          dtype=None, 
          layout=torch.strided, 
          device=None, 
          requires_grad=False)

Parameters

Argument	Description
n (int)	It defines the number of rows in the output tensor.
m (int, optional)	It defines the number of columns for the output tensor, with the default being n.
out (Tensor, optional)	If there is a pre-allocated tensor, you can use this argument to store the result in that tensor. The default is None.
dtype (torch.dtype, optional)	It is the desired data type of the output tensor. The default is None.
layout (torch.layout, optional)	By default, it is torch.strided, but you can define a valid layout for the output tensor.
device (torch.device, optional)	It represents the device, either CPU or CUDA. By default, it takes the current device on which it is installed.
requires_grad (bool, optional)	If True, it tracks gradients. By default, it is False.

Creating an Identity Matrix

Let’s create a 3 x 3 identity matrix with 1s on the diagonal and 0s elsewhere.

import torch

identity_matrix = torch.eye(3)

print(identity_matrix)

# Output:
# tensor([[1., 0., 0.],
#         [0., 1., 0.],
#         [0., 0., 1.]])

The above output shows that indices [0, 0], [1, 1], and [2, 2] have a value of 1, and the others have a value of 0.

Rectangular Identity Matrix

Let’s specify n and m to create an n x m matrix with ones on the diagonal up to min(n, m).

import torch

rectangular_matrix = torch.eye(4, 2)

print(rectangular_matrix)

# Output:
# tensor([[1., 0.],
#         [0., 1.],
#         [0., 0.],
#         [0., 0.]])

You can see that in the above code, the diagonal stops at min(4, 2)=2. So, there are only two 1s in the matrix.

Specifying data Type and device

You can use dtype and device arguments to control the output tensor’s properties.

import torch

eye_cuda = torch.eye(2, dtype=torch.float16, device='cuda')

print(eye_cuda)

Gradient Tracking

If you are working on neural network training, you can enable requires_grad=True.

import torch

eye_grad = torch.eye(4, requires_grad=True)

print(eye_grad)

# Output:
# tensor([[1., 0., 0., 0.],
#         [0., 1., 0., 0.],
#         [0., 0., 1., 0.],
#         [0., 0., 0., 1.]], requires_grad=True)

It can be helpful in backpropagation.

Comparison with torch.diag()

The main difference between torch.eye() and torch.diag() is that .eye() creates a full identity matrix, whereas .diag() can extract or create diagonal tensors.

import torch

eye_identity = torch.eye(2)
print(eye_identity)
# Output:
# tensor([[1., 0.],
#         [0., 1.])

diag_identity = torch.diag(torch.tensor([1., 1.]))
print(diag_identity)
# Output:
# tensor([[1., 0.],
#         [0., 1.])

Here, in the above code, we used a torch.diag([1., 1.]) to create a result similar to torch.eye(2), but torch.diag() is more flexible for custom diagonals. For example, you can pass any value based on the requirement in the case of .diag(), which is not possible in .eye().

That’s all!