torch.nonzero returns wrong indices, and indexing with its result fails at run time #24763

🐛 Describe the bug

torch.nonzero returns wrong indices in ExecuTorch, with no error. Indexing with its result, a common pattern, fails at run time.

import torch
from executorch.exir import to_edge_transform_and_lower
from executorch.runtime import Runtime


class M(torch.nn.Module):
    def forward(self, x):
        return torch.nonzero(x)


x = torch.tensor([[0.0, 1.0, 2.0], [3.0, 0.0, 0.0]])
pte = to_edge_transform_and_lower(torch.export.export(M(), (x,))).to_executorch()
print(Runtime.get().load_program(pte.buffer).load_method("forward").execute([x])[0].tolist())
# ExecuTorch: [[0, 2], [1, 1], [0, 0]]
# eager:      [[0, 1], [0, 2], [1, 0]]

If the model instead returns x[idx[:, 0], idx[:, 1]] with idx = torch.nonzero(x), every run fails:

Check failed (all_contiguous || all_channels_last): 2 input tensors have different dim orders
KernelCall failed at instruction 0:1 in operator aten::select_copy.int_out: 0x12

It fails the same way with and without the XNNPACK partitioner.

Cause. PyTorch returns the nonzero result in column-major layout. Both the eager and the fake tensor have strides (1, N). So the edge program gives that output the dim order (1, 0). The portable nonzero kernel ignores this and always writes row by row. The runtime then reads row-major data as column-major, which gives the wrong indices above. select_copy requires its input and output to have the same dim order, so (1, 0) against (0,) fails its check.

Possible fixes: give the nonzero output a contiguous dim order in the edge program, or make the kernel write in the dim order of its output.

Versions

• executorch 1.6.0 (PyPI) and the 2026-11-05 nightly
• torch 2.15.0, and 2.16 nightly with the ExecuTorch nightly
• Linux x86_64, CPU