pub struct TensorStd {
pub data: Vec<f32>,
pub shape: Vec<usize>,
}
fn elementwise_view_shape(total_len: usize) -> [usize; 5] {
[1, 1, 1, 1, total_len]
}
impl TensorStd {
pub fn new(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, &'static str> {
let mut n = 1usize;
for &d in &shape {
n = n.saturating_mul(d);
}
if n != data.len() {
return Err("data length does not match shape");
}
Ok(Self { data, shape })
}
pub fn zeros(shape: Vec<usize>) -> Self {
let mut n = 1usize;
for &d in &shape {
n = n.saturating_mul(d);
}
Self {
data: vec![0f32; n],
shape,
}
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn rank(&self) -> usize {
self.shape.len()
}
pub fn dims(&self) -> &[usize] {
&self.shape
}
pub fn idx_linear(&self, idxs: &[usize]) -> Option<usize> {
if idxs.len() != self.shape.len() {
return None;
}
let mut linear = 0usize;
for (i, &ind) in idxs.iter().enumerate() {
let dim = self.shape[i];
if ind >= dim {
return None;
}
linear = linear.checked_mul(dim)?;
linear = linear.checked_add(ind)?;
}
if linear < self.data.len() {
Some(linear)
} else {
None
}
}
pub fn get(&self, idxs: &[usize]) -> Option<f32> {
self.idx_linear(idxs).map(|i| self.data[i])
}
pub fn get_mut(&mut self, idxs: &[usize]) -> Option<&mut f32> {
let i = self.idx_linear(idxs)?;
Some(&mut self.data[i])
}
pub fn as_slice(&self) -> &[f32] {
&self.data
}
pub fn as_mut_slice(&mut self) -> &mut [f32] {
&mut self.data
}
pub fn into_vec(self) -> Vec<f32> {
self.data
}
}
pub fn tensor_fill(t: &mut TensorStd, value: f32) {
t.fill(value);
}
pub fn tensor_scale_in_place(t: &mut TensorStd, scale: f32) {
t.scale_in_place(scale);
}
pub fn tensor_add_in_place(a: &mut TensorStd, b: &TensorStd) -> bool {
a.add_in_place(b)
}
impl TensorStd {
pub fn fill(&mut self, value: f32) {
let len = self.data.len();
let mut view = native_neural_network::tensor::TensorView {
data: self.data.as_mut_slice(),
shape: elementwise_view_shape(len),
};
native_neural_network::tensor::tensor_fill(&mut view, value);
}
pub fn tensor_fill(t: &mut TensorStd, value: f32) {
t.fill(value);
}
pub fn scale_in_place(&mut self, scale: f32) {
let len = self.data.len();
let mut view = native_neural_network::tensor::TensorView {
data: self.data.as_mut_slice(),
shape: elementwise_view_shape(len),
};
native_neural_network::tensor::tensor_scale_in_place(&mut view, scale);
}
pub fn add_in_place(&mut self, other: &TensorStd) -> bool {
if self.shape != other.shape || self.data.len() != other.data.len() {
return false;
}
let dst_len = self.data.len();
let src_len = other.data.len();
let mut dst = native_neural_network::tensor::TensorView {
data: self.data.as_mut_slice(),
shape: elementwise_view_shape(dst_len),
};
let mut src_buf = other.data.clone();
let src = native_neural_network::tensor::TensorView {
data: src_buf.as_mut_slice(),
shape: elementwise_view_shape(src_len),
};
native_neural_network::tensor::tensor_add_in_place(&mut dst, &src)
}
}