How to compute the Heaviside step function for each element in input in PyTorch?

To compute the Heaviside step function for each element in the input tensor, we use the torch.heaviside() method. It accepts two parameters − input and values. It returns a new tensor with a computed heaviside step function.

The value of heaviside function is the same as values if input=0. The value of heaviside is zero if input is less than zero. The value of heaviside is 1 if input is greater than zero. It accepts torch tensors of any dimension. It is also called the unit step function.

Syntax

torch.heaviside(input, values)

Steps

We could use the following steps to compute the Heaviside step function −

• Import the required library. In all the following examples, the required Python library is torch. Make sure you have already installed it.

import torch

• Create two tensors − input and values.

input = torch.randn(3,3)
values = torch.tensor([0.5, 0.3, 0.7])

• Compute the heaviside step function of the above defined tensor using torch.heaviside(input, values). Optionally assign this value to a new variable.

hssf = torch.heaviside(input, values)

• Print the above computed Heaviside step function.

print("Heaviside Step Function:
", hssf)

Example 1

In this Python example, we compute the Heaviside step function of a 1D tensor.

import torch
# define input and values tensors
input = torch.tensor([-1.5, 0, 2.0])
values = torch.tensor([0.5])

# display above defined tensors
print("Input Tensor:
", input)
print("Values Tensor:
", values)

# compute heaviside step function
hssf = torch.heaviside(input, values)
print("Heaviside Step Function:
", hssf)

Output

Input Tensor:
   tensor([-1.5000, 0.0000, 2.0000])
Values Tensor:
   tensor([0.5000])
Heaviside Step Function:
   tensor([0.0000, 0.5000, 1.0000])

Example 2

In this example, we compute the Heaviside step function of a 2D tensor.

import torch
# define input and values tensors
input = torch.tensor([[0.2, 0.0, -0.7, -0.2],
   [0.0, 0.6, 0.6, -0.9],
   [0.0, 0.0, 0.0, 0.4],
   [-1.2, 0.0, 0.8, 0.0]])
values = torch.tensor([0.5,0.3, 0.7, 0.8])

# display above defined tensors
print("Input Tensor:
", input)
print("Values Tensor:
", values)

# compute heaviside step function
hssf = torch.heaviside(input, values)
print("Heaviside Step Function:
", hssf)

Output

Input Tensor:
   tensor([[ 0.2000, 0.0000, -0.7000, -0.2000],
      [ 0.0000, 0.6000, 0.6000, -0.9000],
      [ 0.0000, 0.0000, 0.0000, 0.4000],
      [-1.2000, 0.0000, 0.8000, 0.0000]])
Values Tensor:
   tensor([5.0000, 0.3000, 0.7000, 0.8000])
Heaviside Step Function:
   tensor([[1.0000, 0.3000, 0.0000, 0.0000],
      [5.0000, 1.0000, 1.0000, 0.0000],
      [5.0000, 0.3000, 0.7000, 1.0000],
      [0.0000, 0.3000, 1.0000, 0.8000]])