use crate::{Array, Element};
pub trait IntoArray<T: Element> {
fn into_array(self) -> Array<T>;
fn scalar_value(&self) -> Option<T> {
None
}
}
impl<T: Element> IntoArray<T> for Array<T> {
fn into_array(self) -> Array<T> {
self
}
}
impl<T: Element> IntoArray<T> for &Array<T> {
fn into_array(self) -> Array<T> {
self.clone()
}
}
impl<T: Element> IntoArray<T> for T {
fn scalar_value(&self) -> Option<T> {
Some(*self)
}
fn into_array(self) -> Array<T> {
Array::scalar(self)
}
}
impl<T: Element> IntoArray<T> for Vec<T> {
fn into_array(self) -> Array<T> {
Array::from(self)
}
}
impl<T: Element> IntoArray<T> for &[T] {
fn into_array(self) -> Array<T> {
Array::from(self.to_vec())
}
}
impl<T: Element, const N: usize> IntoArray<T> for [T; N] {
fn into_array(self) -> Array<T> {
Array::from(self.to_vec())
}
}
pub fn broadcast_shapes(a: &[usize], b: &[usize]) -> Option<Vec<usize>> {
let n = a.len().max(b.len());
let mut out = Vec::with_capacity(n);
for i in 0..n {
let x = if i < a.len() { a[a.len() - 1 - i] } else { 1 };
let y = if i < b.len() { b[b.len() - 1 - i] } else { 1 };
if x != y && x != 1 && y != 1 {
return None;
}
out.push(if x == 1 { y } else { x });
}
out.reverse();
Some(out)
}
pub fn broadcast_arrays<T: Element, U: Element>(
a: &Array<T>,
b: &Array<U>,
) -> (Array<T>, Array<U>) {
let s = broadcast_shapes(a.shape(), b.shape()).unwrap_or_else(|| {
panic!(
"operands could not be broadcast together with shapes {:?} {:?}",
a.shape(),
b.shape()
)
});
(a.broadcast_to(&s), b.broadcast_to(&s))
}
#[cfg(test)]
mod tests {
#[test]
fn zero_broadcast() {
assert_eq!(super::broadcast_shapes(&[0, 3], &[1, 3]), Some(vec![0, 3]));
}
#[test]
fn incompatible() {
assert_eq!(super::broadcast_shapes(&[2, 3], &[4]), None);
}
}