use std::collections::{BTreeMap, HashMap};
use std::ops::{Add, Mul};
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DType {
F32,
F64,
I8,
I16,
I32,
}
impl DType {
pub fn size(&self) -> usize {
match self {
DType::F32 => 4,
DType::F64 => 8,
DType::I8 => 1,
DType::I16 => 2,
DType::I32 => 4,
}
}
}
pub type TensorShape = Vec<usize>;
#[derive(Debug, Clone)]
pub struct Tensor {
data: Vec<f32>,
shape: TensorShape,
dtype: DType,
}
impl Tensor {
pub fn zeros(shape: &[usize]) -> Self {
let numel: _ = shape.iter().product();
Self {
data: vec![0.0; numel],
shape: shape.to_vec(),
dtype: DType::F32,
}
}
pub fn ones(shape: &[usize]) -> Self {
let numel: _ = shape.iter().product();
Self {
data: vec![1.0; numel],
shape: shape.to_vec(),
dtype: DType::F32,
}
}
pub fn from_vec(data: Vec<f32>, shape: &[usize]) -> Self {
let expected_numel: usize = shape.iter().product();
assert_eq!(data.len(), expected_numel, "Data length must match shape");
Self {
data,
shape: shape.to_vec(),
dtype: DType::F32,
}
}
pub fn randn(shape: &[usize]) -> Self {
let numel: _ = shape.iter().product();
let data: Vec<f32> = (0..numel)
.map(|_| {
let u1: f32 = rand::random::<f32>().max(1e-10);
let u2: f32 = rand::random();
(-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
})
.collect();
Self {
data,
shape: shape.to_vec(),
dtype: DType::F32,
}
}
pub fn shape(&self) -> &[usize] {
&self.shape
}
pub fn numel(&self) -> usize {
self.data.len()
}
pub fn dtype(&self) -> DType {
self.dtype
}
pub fn set_dtype(&mut self, dtype: DType) {
self.dtype = dtype;
}
pub fn get(&self, indices: &[usize]) -> f32 {
let idx: _ = self.flat_index(indices);
self.data[idx]
}
pub fn set(&mut self, indices: &[usize], value: f32) {
let idx: _ = self.flat_index(indices);
self.data[idx] = value;
}
fn flat_index(&self, indices: &[usize]) -> usize {
assert_eq!(indices.len(), self.shape.len());
let mut idx = 0;
let mut stride = 1;
for i in (0..self.shape.len()).rev() {
assert!(indices[i] < self.shape[i]);
idx += indices[i] * stride;
stride *= self.shape[i];
}
idx
}
pub fn mean(&self) -> f32 {
self.data.iter().sum::<f32>() / self.data.len() as f32
}
pub fn sum(&self) -> f32 {
self.data.iter().sum()
}
pub fn reshape(&self, new_shape: &[usize]) -> Self {
let new_numel: usize = new_shape.iter().product();
assert_eq!(new_numel, self.numel(), "New shape must have same number of elements");
Self {
data: self.data.clone(),
shape: new_shape.to_vec(),
dtype: self.dtype,
}
}
pub fn transpose(&self) -> Self {
assert_eq!(self.shape.len(), 2, "Transpose only supports 2D tensors");
let rows: _ = self.shape[0];
let cols: _ = self.shape[1];
let mut data = vec![0.0; self.numel()];
for i in 0..rows {
for j in 0..cols {
data[j * rows + i] = self.data[i * cols + j];
}
}
Self {
data,
shape: vec![cols, rows],
dtype: self.dtype,
}
}
pub fn add(&self, other: &Tensor) -> Self {
assert_eq!(self.shape, other.shape, "Shapes must match for addition");
let data: Vec<f32> = self.data.iter()
.zip(other.data.iter())
.map(|(a, b)| a + b)
.collect();
Self {
data,
shape: self.shape.clone(),
dtype: self.dtype,
}
}
pub fn add_broadcast(&self, other: &Tensor) -> Self {
assert_eq!(other.shape.len(), 1);
assert_eq!(*self.shape.last().unwrap(), other.shape[0]);
let mut data = self.data.clone();
let inner_size: _ = other.shape[0];
for (i, val) in data.iter_mut().enumerate() {
*val += other.data[i % inner_size];
}
Self {
data,
shape: self.shape.clone(),
dtype: self.dtype,
}
}
pub fn mul(&self, other: &Tensor) -> Self {
assert_eq!(self.shape, other.shape);
let data: Vec<f32> = self.data.iter()
.zip(other.data.iter())
.map(|(a, b)| a * b)
.collect();
Self {
data,
shape: self.shape.clone(),
dtype: self.dtype,
}
}
pub fn scale(&self, scalar: f32) -> Self {
let data: Vec<f32> = self.data.iter().map(|x| x * scalar).collect();
Self {
data,
shape: self.shape.clone(),
dtype: self.dtype,
}
}
pub fn matmul(&self, other: &Tensor) -> Self {
assert_eq!(self.shape.len(), 2);
assert_eq!(other.shape.len(), 2);
assert_eq!(self.shape[1], other.shape[0], "Matrix dimensions must match");
let m: _ = self.shape[0];
let k: _ = self.shape[1];
let n: _ = other.shape[1];
let mut data = vec![0.0; m * n];
for i in 0..m {
for j in 0..n {
let mut sum = 0.0;
for l in 0..k {
sum += self.data[i * k + l] * other.data[l * n + j];
}
data[i * n + j] = sum;
}
}
Self {
data,
shape: vec![m, n],
dtype: self.dtype,
}
}
pub fn mse(&self, other: &Tensor) -> f32 {
assert_eq!(self.numel(), other.numel());
let sum: f32 = self.data.iter()
.zip(other.data.iter())
.map(|(a, b)| (a - b).powi(2))
.sum();
sum / self.numel() as f32
}
pub fn data(&self) -> &[f32] {
&self.data
}
pub fn data_mut(&mut self) -> &mut Vec<f32> {
&mut self.data
}
pub fn clone_data(&self) -> Vec<f32> {
self.data.clone()
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_tensor_zeros() {
let t: _ = Tensor::zeros(&[2, 3]);
assert_eq!(t.numel(), 6);
assert!((t.get(&[0, 0])).abs() < 1e-6);
}
}