use serde::{Deserialize, Serialize};
use crate::error::{Error, Result};
use crate::tensor::{Tensor, checked_shape_len};
use super::kind::Framed;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(into = "FrameSeqRepr", try_from = "FrameSeqRepr")]
pub struct FrameSeq {
frames: usize,
dim: usize,
values: Vec<f32>,
}
impl FrameSeq {
pub fn new(frames: usize, dim: usize, values: Vec<f32>) -> Result<Self> {
if dim == 0 {
return Err(Error::validation("a frame sequence needs dim ≥ 1"));
}
let expected = checked_shape_len(&[frames, dim]).ok_or_else(|| {
Error::validation(format!("frame-seq shape [{frames}, {dim}] overflows usize"))
})?;
if values.len() != expected {
return Err(Error::shape(vec![frames, dim], vec![values.len()]));
}
if !values.iter().all(|v| v.is_finite()) {
return Err(Error::NonFinite {
context: "frame sequence values",
});
}
Ok(Self {
frames,
dim,
values,
})
}
pub fn from_tensor(tensor: Tensor) -> Result<Self> {
match *tensor.shape() {
[frames, dim] => {
let values = tensor.data().to_vec();
Self::new(frames, dim, values)
}
ref shape => Err(Error::shape(vec![0, 0], shape.to_vec())),
}
}
pub fn from_data(data: burn_tensor::TensorData) -> Result<Self> {
Self::from_tensor(Tensor::from_data(data)?)
}
pub fn frames(&self) -> usize {
self.frames
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn values(&self) -> &[f32] {
&self.values
}
pub fn frame(&self, index: usize) -> &[f32] {
&self.values[index * self.dim..(index + 1) * self.dim]
}
pub fn to_tensor(&self) -> Tensor {
Tensor::new([self.frames, self.dim], self.values.clone())
.expect("frame-seq invariants imply a valid tensor")
}
pub fn into_data(self) -> burn_tensor::TensorData {
burn_tensor::TensorData::new(self.values, vec![self.frames, self.dim])
}
}
impl Framed for FrameSeq {
fn frames(&self) -> usize {
self.frames
}
fn slice_frames(&self, start: usize, len: usize) -> Self {
let values = self.values[start * self.dim..(start + len) * self.dim].to_vec();
Self {
frames: len,
dim: self.dim,
values,
}
}
}
#[derive(Serialize, Deserialize)]
struct FrameSeqRepr {
shape: [usize; 2],
values: Vec<f32>,
}
impl From<FrameSeq> for FrameSeqRepr {
fn from(seq: FrameSeq) -> Self {
Self {
shape: [seq.frames, seq.dim],
values: seq.values,
}
}
}
impl TryFrom<FrameSeqRepr> for FrameSeq {
type Error = Error;
fn try_from(repr: FrameSeqRepr) -> Result<Self> {
Self::new(repr.shape[0], repr.shape[1], repr.values)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn invariants_hold_at_every_gate() {
assert!(FrameSeq::new(2, 0, vec![]).is_err());
assert!(matches!(
FrameSeq::new(2, 2, vec![0.0; 3]),
Err(Error::Shape { .. })
));
assert!(matches!(
FrameSeq::new(1, 2, vec![0.0, f32::NAN]),
Err(Error::NonFinite { .. })
));
assert!(FrameSeq::new(0, 3, vec![]).is_ok()); assert!(FrameSeq::new(usize::MAX, 2, vec![]).is_err()); }
#[test]
fn from_tensor_demands_rank_two() {
let rank1 = Tensor::vector([1.0, 2.0]);
assert!(matches!(
FrameSeq::from_tensor(rank1),
Err(Error::Shape { .. })
));
let rank2 = Tensor::new([2, 2], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
let seq = FrameSeq::from_tensor(rank2).unwrap();
assert_eq!((seq.frames(), seq.dim()), (2, 2));
}
#[test]
fn rows_and_slices() {
let seq = FrameSeq::new(3, 2, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
assert_eq!(seq.frame(1), &[2.0, 3.0]);
let mid = seq.slice_frames(1, 2);
assert_eq!((mid.frames(), mid.dim()), (2, 2));
assert_eq!(mid.values(), &[2.0, 3.0, 4.0, 5.0]);
}
#[test]
fn serde_is_human_readable_and_gated() {
let seq = FrameSeq::new(1, 2, vec![0.25, -0.5]).unwrap();
let json = serde_json::to_string(&seq).unwrap();
assert_eq!(json, r#"{"shape":[1,2],"values":[0.25,-0.5]}"#);
assert_eq!(serde_json::from_str::<FrameSeq>(&json).unwrap(), seq);
assert!(serde_json::from_str::<FrameSeq>(r#"{"shape":[2,2],"values":[0.0]}"#).is_err());
assert!(serde_json::from_str::<FrameSeq>(r#"{"shape":[1,1],"values":[null]}"#).is_err());
}
#[test]
fn device_boundary_round_trip() {
let seq = FrameSeq::new(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
let back = FrameSeq::from_data(seq.clone().into_data()).unwrap();
assert_eq!(back, seq);
}
}