| 1719 | |
| 1720 | @derived_from(np) |
| 1721 | def unique(ar, return_index=False, return_inverse=False, return_counts=False): |
| 1722 | # Test whether the downstream library supports structured arrays. If the |
| 1723 | # `np.empty_like` call raises a `TypeError`, the downstream library (e.g., |
| 1724 | # CuPy) doesn't support it. In that case we return the |
| 1725 | # `unique_no_structured_arr` implementation, otherwise (e.g., NumPy) just |
| 1726 | # continue as normal. |
| 1727 | try: |
| 1728 | meta = meta_from_array(ar) |
| 1729 | np.empty_like(meta, dtype=[("a", int), ("b", float)]) |
| 1730 | except TypeError: |
| 1731 | return unique_no_structured_arr( |
| 1732 | ar, |
| 1733 | return_index=return_index, |
| 1734 | return_inverse=return_inverse, |
| 1735 | return_counts=return_counts, |
| 1736 | ) |
| 1737 | |
| 1738 | orig_shape = ar.shape |
| 1739 | ar = ar.ravel() |
| 1740 | |
| 1741 | # Run unique on each chunk and collect results in a Dask Array of |
| 1742 | # unknown size. |
| 1743 | |
| 1744 | args = [ar, "i"] |
| 1745 | out_dtype = [("values", ar.dtype)] |
| 1746 | if return_index: |
| 1747 | args.extend([arange(ar.shape[0], dtype=np.intp, chunks=ar.chunks[0]), "i"]) |
| 1748 | out_dtype.append(("indices", np.intp)) |
| 1749 | else: |
| 1750 | args.extend([None, None]) |
| 1751 | if return_counts: |
| 1752 | args.extend([ones((ar.shape[0],), dtype=np.intp, chunks=ar.chunks[0]), "i"]) |
| 1753 | out_dtype.append(("counts", np.intp)) |
| 1754 | else: |
| 1755 | args.extend([None, None]) |
| 1756 | |
| 1757 | out = blockwise(_unique_internal, "i", *args, dtype=out_dtype, return_inverse=False) |
| 1758 | out._chunks = tuple((np.nan,) * len(c) for c in out.chunks) |
| 1759 | |
| 1760 | # Take the results from the unique chunks and do the following. |
| 1761 | # |
| 1762 | # 1. Collect all results as arguments. |
| 1763 | # 2. Concatenate each result into one big array. |
| 1764 | # 3. Pass all results as arguments to the internal unique again. |
| 1765 | # |
| 1766 | # TODO: This should be replaced with a tree reduction using this strategy. |
| 1767 | # xref: https://github.com/dask/dask/issues/2851 |
| 1768 | |
| 1769 | out_parts = [out["values"]] |
| 1770 | if return_index: |
| 1771 | out_parts.append(out["indices"]) |
| 1772 | else: |
| 1773 | out_parts.append(None) |
| 1774 | if return_counts: |
| 1775 | out_parts.append(out["counts"]) |
| 1776 | else: |
| 1777 | out_parts.append(None) |
| 1778 | |