use crate::tensor::{idx_to_global, Tensor};
use crate::cpu::broadcasting::helpers::{get_broadcast_shape};
impl<T> Tensor<T>
where
T: Default + std::ops::Add<Output = T> + Copy,
{
pub fn tens_broadcast_add(&self, tens2: &Tensor<T>) -> Option<Tensor<T>>{
let broadcast_test = get_broadcast_shape(self.get_shape(), tens2.get_shape());
if broadcast_test.is_none(){
return None;
}
let broadcast_shape = broadcast_test.unwrap();
let output_capacity = broadcast_shape.iter().product::<u32>();
let mut return_data = Vec::with_capacity(output_capacity as usize);
for i in 0..output_capacity{
let output_position = idx_to_global(i, &broadcast_shape);
let self_position: Vec<u32> = output_position.iter().zip(self.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
let tens2_position: Vec<u32> = output_position.iter().zip(tens2.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
return_data.push(*self.value(&self_position).unwrap() + *tens2.value(&tens2_position).unwrap());
}
Tensor::from_data(&return_data, &broadcast_shape)
}
}
impl<T> Tensor<T>
where
T: Default + std::ops::Sub<Output = T> + Copy,
{
pub fn tens_broadcast_sub(&self, tens2: &Tensor<T>) -> Option<Tensor<T>>{
let broadcast_test = get_broadcast_shape(self.get_shape(), tens2.get_shape());
if broadcast_test.is_none(){
return None;
}
let broadcast_shape = broadcast_test.unwrap();
let output_capacity = broadcast_shape.iter().product::<u32>();
let mut return_data = Vec::with_capacity(output_capacity as usize);
for i in 0..output_capacity{
let output_position = idx_to_global(i, &broadcast_shape);
let self_position: Vec<u32> = output_position.iter().zip(self.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
let tens2_position: Vec<u32> = output_position.iter().zip(tens2.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
return_data.push(*self.value(&self_position).unwrap() - *tens2.value(&tens2_position).unwrap());
}
Tensor::from_data(&return_data, &broadcast_shape)
}
}
impl<T> Tensor<T>
where
T: Default + std::ops::Mul<Output = T> + Copy,
{
pub fn tens_broadcast_mul(&self, tens2: &Tensor<T>) -> Option<Tensor<T>>{
let broadcast_test = get_broadcast_shape(self.get_shape(), tens2.get_shape());
if broadcast_test.is_none(){
return None;
}
let broadcast_shape = broadcast_test.unwrap();
let output_capacity = broadcast_shape.iter().product::<u32>();
let mut return_data = Vec::with_capacity(output_capacity as usize);
for i in 0..output_capacity{
let output_position = idx_to_global(i, &broadcast_shape);
let self_position: Vec<u32> = output_position.iter().zip(self.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
let tens2_position: Vec<u32> = output_position.iter().zip(tens2.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
return_data.push(*self.value(&self_position).unwrap() * *tens2.value(&tens2_position).unwrap());
}
Tensor::from_data(&return_data, &broadcast_shape)
}
}
impl<T> Tensor<T>
where
T: Default + std::ops::Div<Output = T> + Copy,
{
pub fn tens_broadcast_div(&self, tens2: &Tensor<T>) -> Option<Tensor<T>>{
let broadcast_test = get_broadcast_shape(self.get_shape(), tens2.get_shape());
if broadcast_test.is_none(){
return None;
}
let broadcast_shape = broadcast_test.unwrap();
let output_capacity = broadcast_shape.iter().product::<u32>();
let mut return_data = Vec::with_capacity(output_capacity as usize);
for i in 0..output_capacity{
let output_position = idx_to_global(i, &broadcast_shape);
let self_position: Vec<u32> = output_position.iter().zip(self.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
let tens2_position: Vec<u32> = output_position.iter().zip(tens2.get_shape().iter())
.map(|(op, is)| op%is)
.collect();
return_data.push(*self.value(&self_position).unwrap() / *tens2.value(&tens2_position).unwrap());
}
Tensor::from_data(&return_data, &broadcast_shape)
}
}