Understanding stride for the output of nonzero

1

Hello everyone!

After hours of debugging I noticed that .stride() behaves (in my opinion) strangely when applied to the output of a .nonzero() operation. Here is a minimal working example of what I am encountering:

import torch


def test_stride():
    print()
    example = torch.tensor([[[False, True],
                             [False, False]]])
    nonzero = torch.nonzero(example).contiguous()

    print(f"Nonzero output: {nonzero}")                     # Nonzero output: tensor([[0, 0, 1]])
    print(f"Nonzero stride: {nonzero.stride()}")            # Nonzero stride: (1, 1)

    equivalent = torch.tensor([[0, 0, 1]]).contiguous()
    print(f"Equivalent output: {equivalent}")               # Equivalent output: tensor([[0, 0, 1]])
    print(f"Equivalent stride: {equivalent.stride()}")      # Equivalent stride: (3, 1)

I would assume that in both cases the stride of the tensor is (3, 1), but no matter what I do to the output tensor of the .nonzero() operation the stride stays (1, 1). This is only the case as long as it contains a single element; as soon as two or more elements are returned the output is as expected.

Is this a bug or am I overlooking something? If the latter is the case please tell me how to resolve my issue. Thanks!

2

This isn’t really a bug, but rather a quirk of how how PyTorch optimizes memory layout for single-element results from nonzero().
The key difference is that:

  • For single elements, PyTorch uses a compact stride of (1, 1)
  • For multiple elements, it uses the expected stride of (3, 1)
    This different stride pattern doesn’t make the operations wrong, but it can be surprising when you’re explicitly working with strides.
    To solve this issue, if you need consistent stride behavior you can convert the single element case to match the multi-element case. I would try:
nonzero = torch.nonzero(example)
if nonzero.shape[0] == 1:
    nonzero = nonzero.clone().view(1, -1)  # This will give you (3, 1) stride
1 Like
3

Thanks for the reply! Although this still seems kind of unintuitive for me, your solution (or in practice one inspired by it) works perfectly fine, thank you!

1 Like