use crate::{AutodiffError, Tensor};
pub trait ArrayData {
fn into_array_data(self) -> Result<(Vec<f32>, Vec<usize>), AutodiffError>;
}
impl ArrayData for Vec<f32> {
fn into_array_data(self) -> Result<(Vec<f32>, Vec<usize>), AutodiffError> {
let n = self.len();
Ok((self, vec![n]))
}
}
impl ArrayData for &[f32] {
fn into_array_data(self) -> Result<(Vec<f32>, Vec<usize>), AutodiffError> {
Ok((self.to_vec(), vec![self.len()]))
}
}
impl<const N: usize> ArrayData for [f32; N] {
fn into_array_data(self) -> Result<(Vec<f32>, Vec<usize>), AutodiffError> {
Ok((self.to_vec(), vec![N]))
}
}
impl ArrayData for Vec<Vec<f32>> {
fn into_array_data(self) -> Result<(Vec<f32>, Vec<usize>), AutodiffError> {
let rows = self.len();
let cols = self.first().map_or(0, |row| row.len());
for (i, row) in self.iter().enumerate() {
if row.len() != cols {
return Err(AutodiffError::InvalidArgument(format!(
"compat::array: jagged 2-D 入力(行 {i} の長さ {} が行 0 の長さ {cols} と不一致)",
row.len()
)));
}
}
let mut data = Vec::with_capacity(rows * cols);
for row in self {
data.extend(row);
}
Ok((data, vec![rows, cols]))
}
}
impl<const M: usize, const N: usize> ArrayData for [[f32; N]; M] {
fn into_array_data(self) -> Result<(Vec<f32>, Vec<usize>), AutodiffError> {
let mut data = Vec::with_capacity(M * N);
for row in self {
data.extend(row);
}
Ok((data, vec![M, N]))
}
}
pub fn array<A: ArrayData>(data: A) -> Result<Tensor<f32>, AutodiffError> {
let (flat, shape) = data.into_array_data()?;
Ok(Tensor::new(flat, &shape)?)
}
#[cfg(test)]
mod tests {
use super::*;
fn dense_vec(t: &Tensor<f32>) -> Vec<f32> {
t.contiguous()
.as_slice()
.expect("contiguous() 直後は必ず as_slice() が Some を返す")
.to_vec()
}
#[test]
fn array_1d_vec_infers_shape() {
let t = array(vec![1.0_f32, 2.0, 3.0]).unwrap();
assert_eq!(t.shape(), &[3]);
assert_eq!(dense_vec(&t), vec![1.0, 2.0, 3.0]);
}
#[test]
fn array_1d_slice_infers_shape() {
let src = [1.0_f32, 2.0, 3.0];
let t = array(src.as_slice()).unwrap();
assert_eq!(t.shape(), &[3]);
assert_eq!(dense_vec(&t), vec![1.0, 2.0, 3.0]);
}
#[test]
fn array_1d_literal_infers_shape() {
let t = array([1.0_f32, 2.0, 3.0]).unwrap();
assert_eq!(t.shape(), &[3]);
assert_eq!(dense_vec(&t), vec![1.0, 2.0, 3.0]);
}
#[test]
fn array_2d_vec_infers_shape_row_major() {
let t = array(vec![vec![1.0_f32, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]]).unwrap();
assert_eq!(t.shape(), &[3, 2]);
assert_eq!(dense_vec(&t), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn array_2d_literal_infers_shape_row_major() {
let t = array([[1.0_f32, 2.0], [3.0, 4.0]]).unwrap();
assert_eq!(t.shape(), &[2, 2]);
assert_eq!(dense_vec(&t), vec![1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn array_2d_jagged_input_is_rejected() {
let err = array(vec![vec![1.0_f32, 2.0], vec![3.0]]).unwrap_err();
assert!(matches!(err, AutodiffError::InvalidArgument(_)));
}
#[test]
fn array_empty_1d_is_valid_zero_size_axis() {
let t = array(Vec::<f32>::new()).unwrap();
assert_eq!(t.shape(), &[0]);
}
}