Add assertions that the destination size is the number of unsorted dims when sorting a size-1 axis

ndgrigorian · ndgrigorian · commit 9246251ec265 · 2025-07-07T16:19:56.000-07:00
diff --git a/dpctl/tensor/libtensor/source/sorting/py_argsort_common.hpp b/dpctl/tensor/libtensor/source/sorting/py_argsort_common.hpp
@@ -147,6 +147,7 @@ py_argsort(const dpctl::tensor::usm_ndarray &src,
             return std::make_pair(keep_args_alive_ev, comp_ev);
         }
         else {
+            assert(dst.get_size() == iter_nelems);
             int dst_elemsize = dst.get_elemsize();
             static constexpr int memset_val(0);
 
diff --git a/dpctl/tensor/libtensor/source/sorting/py_sort_common.hpp b/dpctl/tensor/libtensor/source/sorting/py_sort_common.hpp
@@ -148,6 +148,7 @@ py_sort(const dpctl::tensor::usm_ndarray &src,
             return std::make_pair(keep_args_alive_ev, comp_ev);
         }
         else {
+            assert(dst.get_size() == iter_nelems);
             int src_elemsize = src.get_elemsize();
 
             sycl::event copy_ev =

Original file line number	Diff line number	Diff line change
`@@ -147,6 +147,7 @@ py_argsort(const dpctl::tensor::usm_ndarray &src,`
`147`	`147`	`return std::make_pair(keep_args_alive_ev, comp_ev);`
`148`	`148`	`}`
`149`	`149`	`else {`
	`150`	`+ assert(dst.get_size() == iter_nelems);`
`150`	`151`	`int dst_elemsize = dst.get_elemsize();`
`151`	`152`	`static constexpr int memset_val(0);`
`152`	`153`
Original file line number	Diff line number	Diff line change
`@@ -148,6 +148,7 @@ py_sort(const dpctl::tensor::usm_ndarray &src,`
`148`	`148`	`return std::make_pair(keep_args_alive_ev, comp_ev);`
`149`	`149`	`}`
`150`	`150`	`else {`
	`151`	`+ assert(dst.get_size() == iter_nelems);`
`151`	`152`	`int src_elemsize = src.get_elemsize();`
`152`	`153`
`153`	`154`	`sycl::event copy_ev =`