Merge pull request #2177 from rapidsai/branch-24.02

Forward-merge branch-24.02 to branch-24.04
rapidsai · Feb 12, 2024 · 16cdf90 · 16cdf90
2 parents bf850a9 + dc75590
commit 16cdf90
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/cpp/include/raft/util/input_validation.hpp b/cpp/include/raft/util/input_validation.hpp
@@ -39,8 +39,7 @@ constexpr bool is_row_or_column_major(mdspan<ElementType, Extents, layout_right,
 template <class ElementType, class Extents, class Accessor>
 constexpr bool is_row_or_column_major(mdspan<ElementType, Extents, layout_stride, Accessor> m)
 {
-  return m.stride(0) == typename Extents::index_type(1) ||
-         m.stride(1) == typename Extents::index_type(1);
+  return is_row_major(m) || is_col_major(m);
 }
 
 template <class ElementType, class Extents, class Layout, class Accessor>
@@ -64,7 +63,7 @@ constexpr bool is_row_major(mdspan<ElementType, Extents, layout_right, Accessor>
 template <class ElementType, class Extents, class Accessor>
 constexpr bool is_row_major(mdspan<ElementType, Extents, layout_stride, Accessor> m)
 {
-  return m.stride(1) == typename Extents::index_type(1);
+  return m.stride(1) == typename Extents::index_type(1) && m.stride(0) >= m.extent(1);
 }
 
 template <class ElementType, class Extents, class Layout, class Accessor>
@@ -88,7 +87,7 @@ constexpr bool is_col_major(mdspan<ElementType, Extents, layout_right, Accessor>
 template <class ElementType, class Extents, class Accessor>
 constexpr bool is_col_major(mdspan<ElementType, Extents, layout_stride, Accessor> m)
 {
-  return m.stride(0) == typename Extents::index_type(1);
+  return m.stride(0) == typename Extents::index_type(1) && m.stride(1) >= m.extent(0);
 }
 
 template <class ElementType, class IndexType, size_t... Exts, class Layout, class Accessor>