Add tensor sorting operations #1488
Conversation
laggui
force-pushed
the
feat/tensor-sort-argsort
branch
from
944c1c3 to
12c0f67
Compare
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1488 +/- ##
==========================================
+ Coverage 85.80% 85.86% +0.06%
==========================================
Files 650 659 +9
Lines 72524 74082 +1558
==========================================
+ Hits 62229 63612 +1383
- Misses 10295 10470 +175 ☔ View full report in Codecov by Sentry. |
| pub async fn sort(self, dim: usize) -> Tensor<B, D> { | ||
| Tensor::new(sort::<B, D, Float>(self.primitive, dim, /*descending*/ false).await) | ||
| } | ||
|
|
||
| /// Sort the elements by value in ascending order along a given dimension. | ||
| /// Also returns the indices. | ||
| /// | ||
| /// This sort is unstable (i.e., may reorder equal elements). | ||
| #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))] | ||
| pub async fn sort_with_indices(self, dim: usize) -> (Tensor<B, D>, Tensor<B, D, Int>) { | ||
| check!(TensorCheck::sort_dim::<D>("Sort_with_indices", dim)); | ||
| let (values, indices) = | ||
| sort_with_indices::<B, D, Float>(self.primitive, dim, /*descending*/ false).await; | ||
| (Tensor::new(values), Tensor::new(indices)) | ||
| } | ||
|
|
||
| /// Returns the indices that sort the elements by value in ascending order along a given dimension. | ||
| /// | ||
| /// This sort is unstable (i.e., may reorder equal elements). | ||
| #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))] | ||
| pub async fn argsort(self, dim: usize) -> Tensor<B, D, Int> { |
There was a problem hiding this comment.
Could we put those into numeric to avoid duplication? With newer version of Rust, we can add async in traits.
There was a problem hiding this comment.
Oh I guess I should update my rust version, it wouldn't let me do it initially.
There was a problem hiding this comment.
Nevermind it was a warning not an error 🤦
use of `async fn` in public traits is discouraged as auto trait bounds cannot be specified
There was a problem hiding this comment.
Will leave it like this for now to avoid the warnings. This will be temporary anyways with #1490.
| pub trait ElementComparison { | ||
| /// Returns and [Ordering] between `self` and `other`. | ||
| fn cmp(&self, other: &Self) -> Ordering; | ||
| } |
Checklist
run-checks allscript has been executed.Related Issues/PRs
Closes #1369
Unblocks easy implementation for
topkissue #1421Changes
Added
sort(dim),sort_with_indices(dim)andargsort(dim)numeric tensor ops.Also added
sort_descending(dim),sort_descending_with_indices(dim)andargsort_descending(dim)for sorting in descending order.ElementComparisontrait forElementso elements can return an ordering (as required bysortmethods)tchbackendTODO
run_par,iter_paranditer_range_parmacros defined in burn-ndarrayTesting
Forward and backward unit tests.