ttnn.experimental.sort

ttnn.experimental.sort = Operation(python_fully_qualified_name='ttnn.experimental.sort', function=<ttnn._ttnn.operations.experimental.sort_t object>, preprocess_golden_function_inputs=<function default_preprocess_golden_function_inputs>, golden_function=None, postprocess_golden_function_outputs=<function default_postprocess_golden_function_outputs>, is_cpp_operation=True, is_experimental=False)

Sorts the elements of the input tensor along the specified dimension in ascending order by default. If no dimension is specified, the last dimension of the input tensor is used.

This operation is functionally equivalent to the following PyTorch code:

return torch.sort(input_tensor, dim=-1)

Parameters:

input_tensor (*) – The input tensor to be sorted.

Keyword Arguments:

• dim (*) – The dimension along which to sort. Defaults to -1 (last dimension).

• descending (*) – If True, sorts in descending order. Defaults to False.

• stable (*) – If True, ensures the original order of equal elements is preserved. Defaults to False.

• memory_config (*) – Specifies the memory configuration for the output tensor. Defaults to None.

• out (*) – Preallocated output tensors for the sorted values and indices. Defaults to None.

Example:

import ttnn

# Create a tensor
input_tensor = ttnn.Tensor([3, 1, 2])

# Sort the tensor in ascending order
sorted_tensor, indices = ttnn.experimental.sort(input_tensor)

# Sort the tensor in descending order
sorted_tensor_desc, indices_desc = ttnn.experimental.sort(input_tensor, descending=True)

# Sort along a specific dimension
input_tensor_2d = ttnn.Tensor([[3, 1, 2], [6, 5, 4]])
sorted_tensor_dim, indices_dim = ttnn.experimental.sort(input_tensor_2d, dim=1)