Internal sort/argsort implementation. Args: values: The input values. axis: The axis along which to sort. direction: 'ASCENDING' or 'DESCENDING'. return_argsort: Whether to return the argsort result. Returns: Either the sorted values, or the indices of the sorted values in
(values, axis, direction, return_argsort)
| 110 | |
| 111 | |
| 112 | def _sort_or_argsort(values, axis, direction, return_argsort): |
| 113 | """Internal sort/argsort implementation. |
| 114 | |
| 115 | Args: |
| 116 | values: The input values. |
| 117 | axis: The axis along which to sort. |
| 118 | direction: 'ASCENDING' or 'DESCENDING'. |
| 119 | return_argsort: Whether to return the argsort result. |
| 120 | |
| 121 | Returns: |
| 122 | Either the sorted values, or the indices of the sorted values in the |
| 123 | original tensor. See the `sort` and `argsort` docstrings. |
| 124 | |
| 125 | Raises: |
| 126 | ValueError: If axis is not a constant scalar, or the direction is invalid. |
| 127 | """ |
| 128 | if direction not in _SORT_IMPL: |
| 129 | raise ValueError('%s should be one of %s' % (direction, ', '.join( |
| 130 | sorted(_SORT_IMPL.keys())))) |
| 131 | # Axis must be an integer, not a Tensor. |
| 132 | axis = framework_ops.convert_to_tensor(axis, name='axis') |
| 133 | axis_static = tensor_util.constant_value(axis) |
| 134 | if axis.shape.ndims != 0 or axis_static is None: |
| 135 | raise ValueError('axis must be a constant scalar') |
| 136 | axis_static = int(axis_static) # Avoids NumPy casting error |
| 137 | |
| 138 | values = framework_ops.convert_to_tensor(values, name='values') |
| 139 | |
| 140 | return _SORT_IMPL[direction](values, axis_static, return_argsort) |
| 141 | |
| 142 | |
| 143 | def _descending_sort(values, axis, return_argsort=False): |