@@ -513,6 +513,17 @@ def _factorize(data):
513513 return data
514514
515515
516+ def _permute_dims (data , axes ):
517+ """Helper function to get a suitable permute dims function."""
518+ xp = get_array_namespace (data )
519+ if hasattr (xp , "permute_dims" ):
520+ return xp .permute_dims (data , axes )
521+ elif hasattr (xp , "transpose" ):
522+ return xp .transpose (data , axes )
523+ else :
524+ raise NotImplementedError (f"Unknown transpose method for namespace { xp } " )
525+
526+
516527def nunique (data , axis = None , skipna = True , equalna = True ):
517528 """
518529 Count the number of unique values in this array along the given dimensions
@@ -537,7 +548,7 @@ def nunique(data, axis=None, skipna=True, equalna=True):
537548 # Move axes to be aggregated to the end and stack
538549 new_order = [i for i in range (len (shape )) if i not in axis ] + axis
539550 new_shape = [s for i , s in enumerate (shape ) if i not in axis ] + [- 1 ]
540- data = xp .reshape (xp . permute_dims (data , new_order ), new_shape )
551+ data = xp .reshape (_permute_dims (data , new_order ), new_shape )
541552
542553 if is_duck_dask_array (data ):
543554 unique_counts = _dask_nunique (data )
0 commit comments