Exploring Element-Wise Comparisons in PyTorch Tensors with torch.Tensor.eq()

Purpose

True indicates that the corresponding elements in the input tensors are equal, while False signifies inequality.
Returns a new tensor of the same shape as the input tensors, containing boolean values (True or False).
Compares the elements of two PyTorch tensors on an element-wise basis.

Syntax

torch.eq(input1, input2, out=None)

out (optional): An existing tensor to store the result in. This is useful for avoiding memory allocations when performing multiple comparisons. If not provided, a new tensor is created.
input1 and input2: The two tensors you want to compare. Their shapes must be broadcastable.

Example

import torch

tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([2, 2, 4])

result = torch.eq(tensor1, tensor2)
print(result)  # output: tensor([False, True, False])

Key Points

Data types
The input tensors can have different data types (e.g., torch.float32, torch.int64), but the resulting boolean tensor will always have a data type of torch.bool.
Broadcasting
If the input tensors have different shapes but are broadcastable, the smaller tensor is expanded to match the shape of the larger one. Expansion follows broadcasting rules in PyTorch.
Element-wise comparison
torch.eq() compares corresponding elements at the same index in both tensors.

Advanced Usage

In-place operations
PyTorch offers in-place operations like torch.eq_(input1, input2), which modifies the first input tensor (input1) to contain the element-wise comparison results. However, using in-place operations can sometimes make code harder to read and debug, so use them with caution.
Comparison with a number
You can also compare a tensor with a single number. The number is treated as a constant tensor with the same shape as the input tensor, and the comparison is performed element-wise.

When working with large tensors, consider using operations like torch.all or torch.any to check if all or any elements in a tensor satisfy a condition (e.g., being equal to a specific value).
For checking if two tensors have the same size and elements, use torch.equal(tensor1, tensor2).

Comparing tensors with different shapes (broadcasting)

import torch

tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor(2)

result = torch.eq(tensor1, tensor2)
print(result)  # output: tensor([False, True, False])

In this example, tensor2 is broadcast to match the shape of tensor1. Each element in tensor2 is compared with the corresponding element in tensor1.

Comparing tensor with a number

import torch

tensor = torch.tensor([10, 20, 30])
number = 20

result = torch.eq(tensor, number)
print(result)  # output: tensor([False, True, False])

Here, number is treated as a constant tensor with the same shape as tensor.

Using torch.eq_() for in-place operation (use with caution)

import torch

tensor1 = torch.tensor([5, 7, 9])
tensor2 = torch.tensor([3, 7, 11])

tensor1.eq_(tensor2)
print(tensor1)  # output: tensor([False, True, False])

# Now tensor1 has been modified

This code modifies tensor1 in-place to contain the element-wise comparison results.

Checking if all elements in a tensor are equal to a value

import torch

tensor = torch.tensor([5, 5, 5])
value = 5

all_equal = torch.all(torch.eq(tensor, value))
print(all_equal)  # output: tensor(True)

This approach uses torch.all to check if all elements in tensor satisfy the condition of being equal to value using torch.eq().

import torch

tensor = torch.tensor([3, 7, 1])
value = 7

any_equal = torch.any(torch.eq(tensor, value))
print(any_equal)  # output: tensor(True)

Python equality operator (==)

In most cases, you can use the Python equality operator (==) directly with PyTorch tensors. It behaves similarly to torch.eq(), performing element-wise comparisons and returning a boolean tensor.

Note
However, there are some edge cases where == might not work as expected, such as when one tensor is a scalar and the other is not. It's generally safer to stick with torch.eq() unless you're confident about the tensor shapes and data types.

torch.equal(tensor1, tensor2)

This is a good alternative if you want to check if two tensors have the same size and all their corresponding elements are equal. It returns a single boolean value (True if equal, False otherwise).

Logical operators (torch.lt(), torch.gt(), etc.)

If you need to compare tensors based on inequalities (less than, greater than, etc.), use the appropriate logical operators provided by PyTorch:
- torch.lt(tensor1, tensor2): Less than
- torch.gt(tensor1, tensor2): Greater than
- torch.le(tensor1, tensor2): Less than or equal to
- torch.ge(tensor1, tensor2): Greater than or equal to

torch.where(condition, x, y)

This function allows you to create a new tensor based on a comparison condition.
- condition: A boolean tensor created using torch.eq() or other comparison operations.
- x: The tensor to fill elements where the condition is True.
- y: The tensor to fill elements where the condition is False.

import torch

tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([2, 2, 4])

result = torch.where(torch.eq(tensor1, tensor2), 10, 0)
print(result)  # output: tensor([ 0, 10,  0])

For creating new tensors based on comparison conditions, use torch.where().
For other comparison types (less than, greater than), use the respective logical operators.
If you want to check for overall equality (size and elements), use torch.equal().
If you simply need element-wise equality comparison, torch.eq() or == are suitable.