@@ -17,7 +17,6 @@ use crate::{
1717 Axis ,
1818 Data ,
1919 Dimension ,
20- IntoDimension ,
2120 Zip ,
2221} ;
2322
2625 S : Data < Elem = A > ,
2726 D : Dimension ,
2827 A : Clone + Zero ,
29- D :: Smaller : Copy ,
3028{
3129 /// Upper triangular of an array.
3230 ///
@@ -74,10 +72,12 @@ where
7472
7573 let mut res = Array :: zeros ( self . raw_dim ( ) ) ;
7674 let ncols = self . len_of ( Axis ( n - 1 ) ) ;
77- Zip :: indexed ( self . rows ( ) )
75+ let nrows = self . len_of ( Axis ( n - 2 ) ) ;
76+ let indices = Array :: from_iter ( 0 ..nrows) ;
77+ Zip :: from ( self . rows ( ) )
7878 . and ( res. rows_mut ( ) )
79- . for_each ( |i , src , mut dst| {
80- let row_num = i . into_dimension ( ) . last_elem ( ) ;
79+ . and_broadcast ( & indices )
80+ . for_each ( |src , mut dst , row_num| {
8181 let mut lower = match k >= 0 {
8282 true => row_num. saturating_add ( k as usize ) , // Avoid overflow
8383 false => row_num. saturating_sub ( k. unsigned_abs ( ) ) , // Avoid underflow, go to 0
@@ -135,10 +135,13 @@ where
135135
136136 let mut res = Array :: zeros ( self . raw_dim ( ) ) ;
137137 let ncols = self . len_of ( Axis ( n - 1 ) ) ;
138- Zip :: indexed ( self . rows ( ) )
138+ let nrows = self . len_of ( Axis ( n - 2 ) ) ;
139+ let indices = Array :: from_iter ( 0 ..nrows) ;
140+ Zip :: from ( self . rows ( ) )
139141 . and ( res. rows_mut ( ) )
140- . for_each ( |i, src, mut dst| {
141- let row_num = i. into_dimension ( ) . last_elem ( ) ;
142+ . and_broadcast ( & indices)
143+ . for_each ( |src, mut dst, row_num| {
144+ // let row_num = i.into_dimension().last_elem();
142145 let mut upper = match k >= 0 {
143146 true => row_num. saturating_add ( k as usize ) . saturating_add ( 1 ) , // Avoid overflow
144147 false => row_num. saturating_sub ( ( k + 1 ) . unsigned_abs ( ) ) , // Avoid underflow
0 commit comments