Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Uses broadcast indices to remove D::Smaller: Copy trait bound
  • Loading branch information
akern40 authored and bluss committed Aug 9, 2024
commit 5e5f551b04ca8443e9e9dfd248f4fbc353b22cd9
19 changes: 11 additions & 8 deletions src/tri.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ use crate::{
Axis,
Data,
Dimension,
IntoDimension,
Zip,
};

Expand All @@ -26,7 +25,6 @@ where
S: Data<Elem = A>,
D: Dimension,
A: Clone + Zero,
D::Smaller: Copy,
{
/// Upper triangular of an array.
///
Expand Down Expand Up @@ -74,10 +72,12 @@ where

let mut res = Array::zeros(self.raw_dim());
let ncols = self.len_of(Axis(n - 1));
Zip::indexed(self.rows())
let nrows = self.len_of(Axis(n - 2));
let indices = Array::from_iter(0..nrows);
Zip::from(self.rows())
.and(res.rows_mut())
.for_each(|i, src, mut dst| {
let row_num = i.into_dimension().last_elem();
.and_broadcast(&indices)
.for_each(|src, mut dst, row_num| {
let mut lower = match k >= 0 {
true => row_num.saturating_add(k as usize), // Avoid overflow
false => row_num.saturating_sub(k.unsigned_abs()), // Avoid underflow, go to 0
Expand Down Expand Up @@ -135,10 +135,13 @@ where

let mut res = Array::zeros(self.raw_dim());
let ncols = self.len_of(Axis(n - 1));
Zip::indexed(self.rows())
let nrows = self.len_of(Axis(n - 2));
let indices = Array::from_iter(0..nrows);
Zip::from(self.rows())
.and(res.rows_mut())
.for_each(|i, src, mut dst| {
let row_num = i.into_dimension().last_elem();
.and_broadcast(&indices)
.for_each(|src, mut dst, row_num| {
// let row_num = i.into_dimension().last_elem();
let mut upper = match k >= 0 {
true => row_num.saturating_add(k as usize).saturating_add(1), // Avoid overflow
false => row_num.saturating_sub((k + 1).unsigned_abs()), // Avoid underflow
Expand Down