use crate::TensorError;
use std::borrow::Cow;
#[derive(Clone)]
pub struct Tensor<'data> {
pub(super) shape: Vec<usize>,
data: Cow<'data, [f32]>,
}
impl<'data> Tensor<'data> {
pub fn shape(&self) -> &[usize] {
&self.shape
}
pub fn data(&self) -> &[f32] {
self.data.as_ref()
}
pub fn data_mut(&mut self) -> &mut [f32] {
self.data.to_mut()
}
pub fn zeros(shape: Vec<usize>) -> Self {
let nelement: usize = shape.iter().product();
let data = Cow::Owned(vec![0.0; nelement]);
Self { shape, data }
}
pub fn borrowed(data: &'data [f32], shape: Vec<usize>) -> Result<Self, TensorError> {
let cow: Cow<'data, [f32]> = data.into();
Self::new(cow, shape)
}
pub fn new<T>(data: T, shape: Vec<usize>) -> Result<Self, TensorError>
where
T: Into<Cow<'data, [f32]>>,
{
let data = data.into();
if data.len() != shape.iter().product::<usize>() {
return Err(TensorError::InvalidBuffer {
buffer_size: data.len(),
shape,
});
}
Ok(Self { shape, data })
}
}