Fix Tensor.nonzero as_tuple behavior

CHANGELOG.md

@@ -1,4 +1,5 @@
 ## Main
+-   Fix inverted `as_tuple` behavior in Python `Tensor.nonzero()` (PR #7536) (issue #7166)
 -   Replace OpenMP with oneAPI TBB for all CPU parallelism; Open3D no longer depends on OpenMP. This removes the `libomp` / `libgomp` runtime dependency and the thread oversubscription and crashes caused by loading multiple OpenMP runtimes in one process (e.g. alongside PyTorch in Python). The `WITH_OPENMP` CMake option is removed, oneTBB >= 2021.4.0 is required, and `OMP_NUM_THREADS` is replaced by `open3d.utility.set_max_threads()` (C++: `utility::SetMaxThreads()` or a `tbb::task_arena`). `utility::OMPProgressBar` is removed in favor of the thread-safe `utility::ProgressBar`; `utility::GetThreadNum()` and `utility::InParallel()` are removed (PR #6626) (issues #6196, #6544, #6750)
 -   Add point cloud smoothing algorithms: Moving Least Squares (MLS), Laplacian, Taubin, and bilateral smoothing. These methods provide flexible noise reduction for point clouds with different preservation characteristics (PR #7419).
 -   Add vcpkg support for easier dependency management (PR #7386)

cpp/pybind/core/tensor.cpp

@@ -1100,23 +1100,22 @@ input’s data copied.
             "nonzero",
             [](const Tensor& tensor, bool as_tuple) -> py::object {
                 if (as_tuple) {
-                    return py::cast(tensor.NonZero());
-                } else {
                     return py::cast(tensor.NonZeroNumpy());
+                } else {
+                    return py::cast(tensor.NonZero());
                 }
             },
             "Find the indices of the elements that are non-zero.",
             "as_tuple"_a = false);
     docstring::ClassMethodDocInject(
             m, "Tensor", "nonzero",
             {{"as_tuple",
-              "If ``as_tuple`` is True, returns an int64 tensor of shape "
-              "{num_dims, num_non_zeros}, where the i-th row contains the "
-              "indices of the non-zero elements in i-th dimension of the "
-              "original tensor. If ``as_tuple`` is False, Returns a vector "
-              "of "
-              "int64 Tensors, each containing the indices of the non-zero "
-              "elements in each dimension."}});
+              "If ``as_tuple`` is True, returns a list of int64 Tensors, "
+              "each containing the indices of the non-zero elements in each "
+              "dimension. If ``as_tuple`` is False, returns an int64 tensor "
+              "of shape {num_dims, num_non_zeros}, where the i-th row "
+              "contains the indices of the non-zero elements in i-th "
+              "dimension of the original tensor."}});
     tensor.def("all", &Tensor::All, py::call_guard<py::gil_scoped_release>(),
                py::arg("dim") = py::none(), py::arg("keepdim") = false,
                "Returns true if all elements in the tensor are true. Only "

docs/jupyter/core/tensor.ipynb

@@ -1111,8 +1111,8 @@
    "metadata": {},
    "source": [
     "## Nonzero operations\n",
-    "1. When ```as_tuple``` is ```False```(default), it returns a tensor indices of the elements that are non-zero. Each row in the result contains the indices of a non-zero element in the input. If the input has $n$ dimensions, then the resulting tensor is of size $(z x n)$, where $z$ is the total number of non-zero elements in the input tensor.\n",
-    "2. When ```as_tuple``` is ```True```, it returns a tuple of 1D tensors, one for each dimension in input, each containing the indices of all non-zero elements of input. If the input has $n$ dimension, then the resulting tuple contains $n$ tensors of size $z$, where $z$ is the total number of non-zero elements in the input tensor.\n"
+    "1. When ```as_tuple``` is ```False``` (default), it returns an int64 tensor containing the indices of the non-zero elements. If the input has $n$ dimensions, the result has size $(n x z)$, where row $i$ contains the indices in dimension $i$ and $z$ is the total number of non-zero elements.\n",
+    "2. When ```as_tuple``` is ```True```, it returns a list of 1D int64 tensors, one for each input dimension. The list contains $n$ tensors of size $z$.\n"
    ]
   },
   {
@@ -1131,14 +1131,14 @@
       "Tensor[shape={3, 3}, stride={3, 1}, Int64, CPU:0, 0x55d056510470]\n",
       "\n",
       "a.nonzero() = \n",
-      "[[0 1 2 2]\n",
-      "Tensor[shape={4}, stride={1}, Int64, CPU:0, 0x55d05ed0a290], [0 1 0 1]\n",
-      "Tensor[shape={4}, stride={1}, Int64, CPU:0, 0x55d0bf3f4090]]\n",
-      "\n",
-      "a.nonzero(as_tuple = 1) = \n",
       "[[0 1 2 2],\n",
       " [0 1 0 1]]\n",
-      "Tensor[shape={2, 4}, stride={4, 1}, Int64, CPU:0, 0x55d05758e690]\n"
+      "Tensor[shape={2, 4}, stride={4, 1}, Int64, CPU:0, 0x55d05758e690]\n",
+      "\n",
+      "a.nonzero(as_tuple=True) = \n",
+      "[[0 1 2 2]\n",
+      "Tensor[shape={4}, stride={1}, Int64, CPU:0, 0x55d05ed0a290], [0 1 0 1]\n",
+      "Tensor[shape={4}, stride={1}, Int64, CPU:0, 0x55d0bf3f4090]]\n"
      ]
     }
    ],
@@ -1147,7 +1147,7 @@
     "\n",
     "print(\"a = \\n{}\\n\".format(a))\n",
     "print(\"a.nonzero() = \\n{}\\n\".format(a.nonzero()))\n",
-    "print(\"a.nonzero(as_tuple = 1) = \\n{}\".format(a.nonzero(as_tuple=1)))"
+    "print(\"a.nonzero(as_tuple=True) = \\n{}\".format(a.nonzero(as_tuple=True)))"
    ]
   },
   {

python/test/core/test_core.py

@@ -1068,7 +1068,18 @@ def test_non_zero(device):
     np_x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]])
     np_nonzero_tuple = np.nonzero(np_x)
     o3_x = o3c.Tensor(np_x, device=device)
+
+    expected_tensor = np.vstack(np_nonzero_tuple)
+    o3_nonzero = o3_x.nonzero()
+    o3_nonzero_false = o3_x.nonzero(as_tuple=False)
+    assert isinstance(o3_nonzero, o3c.Tensor)
+    assert isinstance(o3_nonzero_false, o3c.Tensor)
+    np.testing.assert_equal(o3_nonzero.cpu().numpy(), expected_tensor)
+    np.testing.assert_equal(o3_nonzero_false.cpu().numpy(), expected_tensor)
+
     o3_nonzero_tuple = o3_x.nonzero(as_tuple=True)
+    assert not isinstance(o3_nonzero_tuple, o3c.Tensor)
+    assert len(o3_nonzero_tuple) == np_x.ndim
     for np_t, o3_t in zip(np_nonzero_tuple, o3_nonzero_tuple):
         np.testing.assert_equal(np_t, o3_t.cpu().numpy())