Fixes for _unravel_index (#65)

Handle cases where there are multiple labels selected with some structure to them.
dask · Sep 17, 2018 · 3dcf9f3 · 3dcf9f3
1 parent e3ca96e
commit 3dcf9f3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dask_image/ndmeasure/_utils.py b/dask_image/ndmeasure/_utils.py
@@ -70,7 +70,9 @@ def _ravel_shape_indices(dimensions, dtype=int, chunks=None):
 
 
 def _unravel_index_kernel(indices, func_kwargs):
-    return numpy.stack(numpy.unravel_index(indices, **func_kwargs), axis=1)
+    nd_indices = numpy.unravel_index(indices, **func_kwargs)
+    nd_indices = numpy.stack(nd_indices, axis=indices.ndim)
+    return nd_indices
 
 
 def _unravel_index(indices, dims, order='C'):
@@ -85,7 +87,7 @@ def _unravel_index(indices, dims, order='C'):
             _unravel_index_kernel,
             dtype=numpy.intp,
             chunks=indices.chunks + ((len(dims),),),
-            new_axis=1,
+            new_axis=indices.ndim,
             func_kwargs={"dims": dims, "order": order}
         )
     else: