Skip to content

Commit bb7e02f

Browse files
committed
Uses broadcast indices to remove D::Smaller: Copy trait bound
1 parent 5ed317a commit bb7e02f

File tree

1 file changed

+11
-8
lines changed

1 file changed

+11
-8
lines changed

src/tri.rs

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ use crate::{
1717
Axis,
1818
Data,
1919
Dimension,
20-
IntoDimension,
2120
Zip,
2221
};
2322

@@ -26,7 +25,6 @@ where
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

Comments
 (0)