torch.sqrt(): Square Root of Elements in a Tensor

The torch.sqrt() method calculates the element-wise square root of the input tensor. It returns the output tensor containing the square root of each element, with the same shape and data type as the input.

The basic mathematics formula is this: output = √input

The primary applications of the square root in mathematics include calculating Root Mean Square Error (RMSE), distance calculations, normalization, Euclidean norms, and signal processing.

It supports float16, float32, float64, complex64, complex128, and non-negative integers.

The torch.square() is the inverse of the torch.sqrt() operation.

Syntax

torch.sqrt(input, out=None)

Parameters

Argument	Description
input (Tensor)	It represents an input tensor.
out (Tensor, optional)	It represents an output tensor.

Square root of 1D tensor

import torch

tensor_1d = torch.tensor([0.0, 1.0, 4.0, 9.0])

print(tensor_1d)
# Output: tensor([0.0000, 1.0000, 4.0000, 9.0000])

sqrt_tensor = torch.sqrt(tensor_1d)

print(sqrt_tensor)
# Output: tensor([0.0000, 1.0000, 2.0000, 3.0000])

2D tensor

import torch

tensor_2d = torch.tensor([[0.0, 1.0], [4.0, 9.0]])

print(tensor_2d)
# Output: tensor([[0., 1.],
#                 [4., 9.]])

sqrt_2d_tensor = torch.sqrt(tensor_2d)

print(sqrt_2d_tensor)
# Output: tensor([[0., 1.],
#                [2., 3.]])

Negative elements

If the input value is negative, the sqrt() method returns NaN. For real-valued inputs, elements must be ≥ 0.

import torch

negative_tensor = torch.tensor([-4.0, -16.0, -9.0])

print(negative_tensor)
# Output: tensor([ -4., -16.,  -9.])

sqrt_of_negative = torch.sqrt(negative_tensor)

print(sqrt_of_negative)
# Output: tensor([nan, nan, nan])

Integer tensor (Auto-Casts to Float)

If you pass an integer tensor, it will automatically be cast to float.

import torch

int_tensor = torch.tensor([4, 16, 9], dtype=torch.int32)

print(int_tensor)
# Output: tensor([ 4, 16,  9], dtype=torch.int32)

print(int_tensor.dtype)
# Output: torch.int32

sqrt_float = torch.sqrt(int_tensor)

print(sqrt_float)
# Output: tensor([2., 4., 3.])

print(sqrt_float.dtype)
# Output: torch.float32

You can see that the output tensor’s type is torch.float.32.

Complex numbers

If you pass a complex-valued tensor to the sqrt() method, even if it includes negative real values, PyTorch will not return NaN. Instead, it computes the mathematically correct complex square root.

import torch

complex_tensor = torch.tensor([-4.0 + 0j, 3.0 - 4j])

print(complex_tensor)
# Output: tensor([-4.+0.j,  3.-4.j])

print(complex_tensor.dtype)
# Output: torch.complex64

complex_sqrt = torch.sqrt(complex_tensor)

print(complex_sqrt)
# Output: tensor([0.0000+2.0000j, 2.0000-1.0000j])

print(complex_sqrt.dtype)
# Output: torch.complex64

Using the “out” argument

If you have a preallocated tensor, you can store the result of the sqrt() method in this tensor.

import torch

input_tensor = torch.tensor([1.0, 4.0, 16.0])

out = torch.empty(3)

print(out)
# Output: tensor([0., 0., 0.])

torch.sqrt(input_tensor, out=out)

print(out)
# Output: tensor([1., 2., 4.])

In the above code, we created an empty preallocated tensor using torch.empty() method and then store the square root values into this tensor using the “out” parameter. That’s all!