@@ -17,7 +17,6 @@ use crate::{
17
17
Axis ,
18
18
Data ,
19
19
Dimension ,
20
- IntoDimension ,
21
20
Zip ,
22
21
} ;
23
22
26
25
S : Data < Elem = A > ,
27
26
D : Dimension ,
28
27
A : Clone + Zero ,
29
- D :: Smaller : Copy ,
30
28
{
31
29
/// Upper triangular of an array.
32
30
///
@@ -74,10 +72,12 @@ where
74
72
75
73
let mut res = Array :: zeros ( self . raw_dim ( ) ) ;
76
74
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 ( ) )
78
78
. 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| {
81
81
let mut lower = match k >= 0 {
82
82
true => row_num. saturating_add ( k as usize ) , // Avoid overflow
83
83
false => row_num. saturating_sub ( k. unsigned_abs ( ) ) , // Avoid underflow, go to 0
@@ -135,10 +135,13 @@ where
135
135
136
136
let mut res = Array :: zeros ( self . raw_dim ( ) ) ;
137
137
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 ( ) )
139
141
. 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();
142
145
let mut upper = match k >= 0 {
143
146
true => row_num. saturating_add ( k as usize ) . saturating_add ( 1 ) , // Avoid overflow
144
147
false => row_num. saturating_sub ( ( k + 1 ) . unsigned_abs ( ) ) , // Avoid underflow
0 commit comments