The NumPy Square Root Operation

Ben Cook • Posted 2020-12-30 • Last updated 2021-10-15
by

In NumPy, you can find the square root of a number of sequence with np.sqrt():

import numpy as np

np.sqrt(4)

# Expected result
# 2.0

np.sqrt([4, 16, 64])

# Expected result
# array([2., 4., 8.])

The input is “array-like”, meaning it can be a number, a Python sequence or a NumPy array (with any number of dimensions). The result will be a NumPy scalar (if you pass in a scalar) or a NumPy array (if you pass in a sequence or an array).

np.sqrt() is an element-wise function, meaning it operates on each element of an array independently. The output will be the same shape as the input and each of the output elements will be the square root of its corresponding input element. Easy!

A few other fun facts about np.sqrt():

  • Negative numbers will be np.nan (you will also see a RuntimeWarning)
  • The square root of np.inf is np.inf
  • If you pass in integers, NumPy will coerce the result to one of its float data types. As of version 1.19.4, it seems to convert int types to float with double the precision up to float64. So int8 becomes float16, int16 becomes float32 and int32 and above become float64.
  • Complex numbers work as expected (e.g. np.sqrt(8 - 6j)). The result will be one of the complex NumPy types.

By the way, other popular tensor libraries like PyTorch and TensorFlow follow roughly the same API. See the documentation for torch.sqrt() and tf.sqrt().