# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=import-outside-toplevel, redefined-builtin, unused-argument """Set operators.""" import numpy as np # type: ignore import tvm from ..expr import Expr, prim_value from . import _ffi_api def unique( x: Expr, sorted: bool | Expr = True, return_index: bool | Expr = False, return_inverse: bool | Expr = False, return_counts: bool | Expr = False, axis: int | Expr | None = None, ) -> Expr: """Find the unique elements in a given tensor. In addition, it optionally returns - the indices of the input tensor that give the unique values; - the indices of the unique tensor that reconstruct the input tensor; - the number of times each unique value comes up in the input tensor. Parameters ---------- x : relax.Expr The input tensor. sorted : Union[bool, Expr] Whether to sort the unique elements in ascending order before returning as output. return_index : Union[bool, Expr] Whether to return an additional tensor with indices for where elements in the unique tensor come from the original input. return_inverse : Union[bool, Expr] Whether to return an additional tensor with indices for where elements in the original input ended up in the returned unique list. return_counts : Union[bool, Expr] Whether to return an additional tensor with counts of each unique elements. axis : Optional The dimension to apply unique. If not specified, the unique values of the flattened input are returned. Returns ------- ret : relax.Expr The created relax call with """ if isinstance(sorted, bool): sorted = prim_value(sorted) if isinstance(return_index, bool): return_index = prim_value(return_index) if isinstance(return_inverse, bool): return_inverse = prim_value(return_inverse) if isinstance(return_counts, bool): return_counts = prim_value(return_counts) if axis is not None and isinstance(axis, int): axis = prim_value(axis) return _ffi_api.unique( # type: ignore x, sorted, return_index, return_inverse, return_counts, axis ) @tvm.register_global_func("relax.run.unique") def numpy_unique( x: tvm.runtime.tensor, sorted: int, return_index: int, return_inverse: int, return_counts: int, axis: int | None = None, ) -> tvm.runtime.tensor: """Returns the unique elements of the input tensor. Uses numpy.unique to compute unique elements. """ import builtins x_numpy = x.numpy() # Call numpy.unique with all the requested return flags result = np.unique( x_numpy, return_index=bool(return_index), return_inverse=bool(return_inverse), return_counts=bool(return_counts), axis=axis, ) # If no optional outputs requested, result is just the unique values if not bool(return_index) and not bool(return_inverse) and not bool(return_counts): unique_values = result if not sorted: indices = np.unique(x_numpy, return_index=True, axis=axis)[1] unique_values = np.take(x_numpy, builtins.sorted(indices), axis=axis) return tvm.runtime.tensor(unique_values) # Otherwise, numpy returns a tuple unique_values = result[0] output_list = [] result_idx = 1 # Handle sorting for unique values if not sorted and bool(return_index): # Get the indices from numpy result indices = result[result_idx] result_idx += 1 # Sort indices to get original order sort_order = np.argsort(indices) unique_values = np.take(unique_values, sort_order, axis=axis) indices = np.sort(indices) output_list.append(tvm.runtime.tensor(unique_values)) output_list.append(tvm.runtime.tensor(indices)) elif not sorted: # Need to get indices to reorder _, indices = np.unique(x_numpy, return_index=True, axis=axis) sort_order = np.argsort(indices) unique_values = np.take(unique_values, sort_order, axis=axis) output_list.append(tvm.runtime.tensor(unique_values)) if bool(return_index): indices_from_result = result[result_idx] result_idx += 1 output_list.append(tvm.runtime.tensor(np.sort(indices_from_result))) else: # Sorted case output_list.append(tvm.runtime.tensor(unique_values)) if bool(return_index): output_list.append(tvm.runtime.tensor(result[result_idx])) result_idx += 1 if bool(return_inverse): inverse_indices = result[result_idx] if not sorted: # Need to remap inverse indices to match reordered unique values _, orig_indices = np.unique(x_numpy, return_index=True, axis=axis) sort_order = np.argsort(orig_indices) inverse_mapping = np.empty_like(sort_order) inverse_mapping[sort_order] = np.arange(len(sort_order)) inverse_indices = inverse_mapping[inverse_indices] # ONNX spec: inverse_indices is always 1D # When axis is None, it has length X.size (flattened) # When axis is specified, it has length X.shape[axis] # numpy.unique already returns 1D inverse_indices, so no reshaping needed output_list.append(tvm.runtime.tensor(inverse_indices)) result_idx += 1 if bool(return_counts): counts = result[result_idx] if not sorted: # Reorder counts to match reordered unique values _, orig_indices = np.unique(x_numpy, return_index=True, axis=axis) sort_order = np.argsort(orig_indices) counts = counts[sort_order] output_list.append(tvm.runtime.tensor(counts)) return tuple(output_list) def nonzero(x: Expr) -> Expr: """Find the indices of elements of a tensor that are non-zero. Parameters ---------- x : relax.Expr The input data tensor. Returns ------- result : relax.Expr A 2-D tensor containing indices of non-zero elements. Note ---- This function is equivalent to `onnx.nonzero`. Examples -------- .. code-block:: python x = [[0, 1], [2, 0]] nonzero(x) = [[0, 1], [1, 0]] """ return _ffi_api.nonzero(x) # type: ignore @tvm.register_global_func("relax.run.nonzero") def numpy_nonzero(x: tvm.runtime.tensor) -> tvm.runtime.tensor: np_result = np.atleast_1d(x.numpy()).nonzero() return tvm.runtime.tensor(np.stack(np_result, axis=0))