torch.nansum(): Sum of all non-NaN Elements in a Tensor

The torch.nansum() method calculates the sum of all non-NaN (Not a Number) elements in a tensor, ignoring any NaN values. Basically, if it encounters NaNs, it treats them as zero.

Syntax

torch.nansum(input, dim=None, keepdim=False, dtype=None)

Parameters

Argument	Description
input (Tensor)	It is an input tensor that contains elements, including NaNs.
dim (int or tuple of ints, optional)	It is the dimension(s) along which to compute the sum.
keepdim (bool, optional)	If True, it retains the reduced dimensions with size 1.
dtype (torch.dtype, optional)	It defines the data type of the output tensor.

Summing all elements in a tensor

If you don’t pass any specific dimension, it will try to sum all elements in a tensor.

import torch

tensor = torch.tensor([21.0, float('nan'), 19.0, float('nan'), 5.0])

nansum_all = torch.nansum(tensor)

print(nansum_all)

# Output: tensor(45.)

You can see that it completely ignores NaN values, or you can say it assumes them as 0s while adding all the elements. So, it looks like 21.0 + 0.0 + 19.0 + 0.0 + 5.0 = 45.0.

Summing along a specific dimension

If we are working with either 2D or 3D tensors, we can calculate the sum along a specific dimension, like dim=0 or dim=1.

Calculating the sum along columns means the calculation will be done row-wise. So, the sum of the first row, and then the sum of the second row.

import torch

tensor = torch.tensor([[11.0, float('nan'), 31.0],
                      [41.0, 51.0, float('nan')]])

nansum_along_columns = torch.nansum(tensor, dim=1)

print(nansum_along_columns)

# Output: tensor([42., 92.])

Explanation: For each row:

1. Row 1: [11.0, NaN, 31.0] → 11.0 + 31.0 = 42.0
2. Row 2: [41.0, 51.0, NaN] → 41.0 + 51.0 = 92.0

dim=0

import torch

tensor = torch.tensor([[11.0, float('nan'), 31.0],
                      [41.0, 51.0, float('nan')]])

nansum_along_rows = torch.nansum(tensor, dim=0)

print(nansum_along_rows)

# Output: tensor([52., 51., 31.])

Explanation: For each column:

1. Column 1: [11.0, 41.0] → 11.0 + 41.0 = 52.0
2. Column 2: [nan, 51.0] → 0.0 + 51.0 = 51.0
3. Column3: [31.0, nan] → 31.0 + 0.0 = 31.0

Using keepdim=True

import torch

tensor = torch.tensor([[11.0, float('nan'), 31.0],
                      [41.0, 51.0, float('nan')]])

# Sum along dimension 0 (rows) with keepdim=True
nan_keep = torch.nansum(tensor, dim=0, keepdim=True)

print(nan_keep)

# Output: tensor([[52., 51., 31.]])

Summing along dim=0 (column-wise):

1. Column 1: 11.0 + 41.0 = 52.0
2. Column 2: NaN + 51.0 = 51.0
3. Column 3: 31.0 + NaN = 31.0

The output shape is (1, 3) because keepdim=True is used.

What if the tensor only contains NaNs?

What if you have a tensor that only contains NaN values? What will be the output? Well, the output will be 0.0 because there is nothing to sum.

import torch

tensor = torch.tensor([float('nan'), float('nan'), float('nan')])

only_nans = torch.nansum(tensor)

print(only_nans)

# Output: tensor(0.)

That’s all!