use crate::{prelude::Number, shape::Shape, Alloc, Buffer, Dim1, Dim2};
pub trait WithShape<D, C> {
#[cfg_attr(feature = "cpu", doc = "```")]
#[cfg_attr(not(feature = "cpu"), doc = "```ignore")]
fn with(device: D, array: C) -> Self;
}
impl<'a, T, D, const N: usize> WithShape<&'a D, [T; N]> for Buffer<'a, T, D, Dim1<N>>
where
T: Number, D: Alloc<'a, T, Dim1<N>>,
{
#[inline]
fn with(device: &'a D, array: [T; N]) -> Self {
Buffer::from_array(device, array)
}
}
impl<'a, T, D, const N: usize> WithShape<&'a D, &[T; N]> for Buffer<'a, T, D, Dim1<N>>
where
T: Number,
D: Alloc<'a, T, Dim1<N>>,
{
#[inline]
fn with(device: &'a D, array: &[T; N]) -> Self {
Buffer::from_array(device, *array)
}
}
impl<'a, T, D, const B: usize, const A: usize> WithShape<&'a D, [[T; A]; B]>
for Buffer<'a, T, D, Dim2<B, A>>
where
T: Number,
D: Alloc<'a, T, Dim2<B, A>>,
{
#[inline]
fn with(device: &'a D, array: [[T; A]; B]) -> Self {
Buffer::from_array(device, array)
}
}
impl<'a, T, D, const B: usize, const A: usize> WithShape<&'a D, &[[T; A]; B]>
for Buffer<'a, T, D, Dim2<B, A>>
where
T: Number,
D: Alloc<'a, T, Dim2<B, A>>,
{
#[inline]
fn with(device: &'a D, array: &[[T; A]; B]) -> Self {
Buffer::from_array(device, *array)
}
}
impl<'a, T, D, S: Shape> WithShape<&'a D, ()> for Buffer<'a, T, D, S>
where
D: Alloc<'a, T, S>,
{
fn with(device: &'a D, _: ()) -> Self {
Buffer::new(device, S::LEN)
}
}
#[cfg(test)]
mod tests {
#[cfg(feature = "cpu")]
#[test]
fn test_with_const_dim2_cpu() {
use crate::{Buffer, WithShape, CPU};
let device = CPU::new();
let buf = Buffer::with(&device, [[1.0, 2.0], [3.0, 4.0]]);
assert_eq!(&*buf, &[1.0, 2.0, 3.0, 4.0]);
}
#[cfg(feature = "stack")]
#[test]
fn test_with_const_dim2_stack() {
use crate::{Buffer, Stack, WithShape};
let device = Stack;
let buf = Buffer::with(&device, [[1.0, 2.0], [3.0, 4.0]]);
assert_eq!(buf.ptr.array, [[1.0, 2.0,], [3.0, 4.0]]);
}
}