Reduces a vector of arrays sharing the same first dimension, over the first dimension, using log number of `combine` calls. `combine` receives two vectors (first, second) of same shape, and needs to return one such vector of same shape (the same vector size, and array shapes). `combine` is assumed to be associative and commutative.
(
mut nodes: Vec<Node>,
combine: impl Fn(Vec<Node>, Vec<Node>) -> Result<Vec<Node>>,
)
| 361 | /// vector of same shape (the same vector size, and array shapes). |
| 362 | /// `combine` is assumed to be associative and commutative. |
| 363 | pub fn custom_reduce_vec( |
| 364 | mut nodes: Vec<Node>, |
| 365 | combine: impl Fn(Vec<Node>, Vec<Node>) -> Result<Vec<Node>>, |
| 366 | ) -> Result<Vec<Node>> { |
| 367 | if nodes.is_empty() { |
| 368 | return Err(runtime_error!("Can't reduce an empty vector")); |
| 369 | } |
| 370 | let ns: Vec<u64> = nodes |
| 371 | .iter() |
| 372 | .map(|node| Ok(node.get_type()?.get_dimensions()[0])) |
| 373 | .collect::<Result<_>>()?; |
| 374 | let mut n = ns[0]; |
| 375 | if ns.iter().any(|el| *el != n) { |
| 376 | return Err(runtime_error!("All nodes must share the first dimension")); |
| 377 | } |
| 378 | |
| 379 | let mut result = None; |
| 380 | while n > 0 { |
| 381 | if n % 2 == 1 { |
| 382 | let (first, rest) = nodes |
| 383 | .into_iter() |
| 384 | .map(|node| { |
| 385 | Ok(( |
| 386 | node.get(vec![0])?, |
| 387 | if n > 1 { |
| 388 | node.get_slice(vec![SliceElement::SubArray(Some(1), None, None)])? |
| 389 | } else { |
| 390 | // There is nothing left in the array when we remove the first row. |
| 391 | // get_slice() would then give an error, but we're done anyway, so we |
| 392 | // can assign anything here. |
| 393 | node |
| 394 | }, |
| 395 | )) |
| 396 | }) |
| 397 | .collect::<Result<Vec<(Node, Node)>>>()? |
| 398 | .into_iter() |
| 399 | .unzip(); |
| 400 | result = match result { |
| 401 | None => Some(first), |
| 402 | Some(result) => Some(combine(result, first)?), |
| 403 | }; |
| 404 | nodes = rest; |
| 405 | n -= 1; |
| 406 | } else { |
| 407 | let (half1, half2) = nodes |
| 408 | .into_iter() |
| 409 | .map(|node| { |
| 410 | Ok(( |
| 411 | node.get_slice(vec![SliceElement::SubArray( |
| 412 | Some(0), |
| 413 | Some((n / 2) as i64), |
| 414 | None, |
| 415 | )])?, |
| 416 | node.get_slice(vec![SliceElement::SubArray( |
| 417 | Some((n / 2) as i64), |
| 418 | None, |
| 419 | None, |
| 420 | )])?, |
no test coverage detected