helena 0.1.0

Core types and component interfaces for helena, a latent data-to-waveform generation platform.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240

//! The host-side tensor: a thin wrapper over Burn's [`TensorData`].
//!
//! helena's *boundary* representation for tensors is host-side and
//! serde-friendly, so latents can be exported, hashed, and disk-cached
//! without a backend or device. Device-side math runs on
//! `burn::tensor::Tensor<B, D>`; convert across the boundary with
//! [`Tensor::into_data`] / [`Tensor::from_data`].
//!
//! `Tensor` is the *raw carrier* (A2): it proves dtype (`f32`) and
//! shape-length consistency, nothing more. Domain meaning — rank, finiteness,
//! axis semantics — belongs to the domain newtypes that wrap it
//! ([`FrameSeq`](crate::latent::FrameSeq), [`Pcm`](crate::signal::Pcm), the
//! conditioning slots). Where data is genuinely heterogeneous (source
//! features), the raw carrier is the honest type.

use burn_tensor::TensorData;
use serde::{Deserialize, Serialize};

use crate::error::Error;

/// A host-side, serde-serializable `f32` tensor (a wrapper over [`TensorData`]).
///
/// Invariant, enforced by every constructor and re-checked on
/// deserialization: the payload is `f32` and holds exactly `product(shape)`
/// elements, a product that does not overflow `usize`. The wrapped data is
/// private and there is no unchecked constructor, so
/// [`shape`](Self::shape), [`len`](Self::len), and [`data`](Self::data) can
/// never disagree. Deserialization routes through [`TryFrom`], so a
/// hand-built record with the wrong dtype or a ragged shape fails closed.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(into = "TensorData", try_from = "TensorData")]
pub struct Tensor(TensorData);

/// Row-major element count for `shape`, or `None` if the product overflows
/// `usize`. Shared with the domain newtypes so the overflow guard against
/// crafted shapes lives in exactly one place: a wrapped product could
/// otherwise pass a length check and declare a huge tensor backed by tiny
/// data.
pub(crate) fn checked_shape_len(shape: &[usize]) -> Option<usize> {
    shape
        .iter()
        .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
}

impl Tensor {
    /// Build a tensor from a shape and its row-major `f32` data.
    ///
    /// Returns [`Error::Shape`] if the data length does not equal the product
    /// of the shape dimensions, or [`Error::Validation`] if that product
    /// overflows `usize` (a crafted shape like `[usize::MAX, 2]` would
    /// otherwise wrap to a small product and pass the length check).
    pub fn new(shape: impl Into<Vec<usize>>, data: impl Into<Vec<f32>>) -> Result<Self, Error> {
        let shape = shape.into();
        let data = data.into();
        let expected = checked_shape_len(&shape)
            .ok_or_else(|| Error::validation(format!("tensor shape {shape:?} overflows usize")))?;
        if expected != data.len() {
            return Err(Error::shape(shape, vec![data.len()]));
        }
        Ok(Self(TensorData::new(data, shape)))
    }

    /// A zero-filled tensor of the given shape.
    ///
    /// Returns [`Error::Validation`] if the shape's element count overflows
    /// `usize` — the same guard [`new`](Self::new) applies.
    pub fn zeros(shape: impl Into<Vec<usize>>) -> Result<Self, Error> {
        let shape = shape.into();
        let len = checked_shape_len(&shape)
            .ok_or_else(|| Error::validation(format!("tensor shape {shape:?} overflows usize")))?;
        Ok(Self(TensorData::new(vec![0.0f32; len], shape)))
    }

    /// A 1-D tensor wrapping a vector of values.
    pub fn vector(data: impl Into<Vec<f32>>) -> Self {
        let data = data.into();
        let len = data.len();
        Self(TensorData::new(data, vec![len]))
    }

    /// The tensor's shape.
    pub fn shape(&self) -> &[usize] {
        &self.0.shape
    }

    /// The flat, row-major `f32` data.
    pub fn data(&self) -> &[f32] {
        self.0
            .as_slice::<f32>()
            .expect("helena tensors always store f32")
    }

    /// Number of elements.
    pub fn len(&self) -> usize {
        self.shape().iter().product()
    }

    /// Whether the tensor has no elements.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Whether every element is finite (no NaN, no infinities).
    pub fn is_finite(&self) -> bool {
        self.data().iter().all(|x| x.is_finite())
    }

    /// Borrow the underlying [`TensorData`].
    pub fn as_data(&self) -> &TensorData {
        &self.0
    }

    /// Consume into [`TensorData`] — the input to
    /// `burn::tensor::Tensor::from_data` when moving onto a device.
    pub fn into_data(self) -> TensorData {
        self.0
    }

    /// Wrap [`TensorData`] coming back from a device tensor's `into_data()`.
    ///
    /// Returns [`Error::Validation`] if the payload is not `f32` or the
    /// shape's element count overflows `usize`, and [`Error::Shape`] if the
    /// element count does not match the shape.
    pub fn from_data(data: TensorData) -> Result<Self, Error> {
        Self::try_from(data)
    }
}

/// Validating gate from raw [`TensorData`] into a [`Tensor`]: the single
/// checked path every untrusted entry point (deserialization,
/// [`Tensor::from_data`]) routes through.
impl TryFrom<TensorData> for Tensor {
    type Error = Error;

    fn try_from(data: TensorData) -> Result<Self, Error> {
        // `as_slice::<f32>()` doubles as the dtype check: it errors unless
        // the payload really is `f32`, and its length is the true element
        // count.
        let actual = data
            .as_slice::<f32>()
            .map_err(|_| Error::validation("helena tensors must store f32 data"))?
            .len();
        let expected = checked_shape_len(&data.shape).ok_or_else(|| {
            Error::validation(format!("tensor shape {:?} overflows usize", data.shape))
        })?;
        if actual != expected {
            return Err(Error::shape(data.shape.clone(), vec![actual]));
        }
        Ok(Self(data))
    }
}

impl From<Tensor> for TensorData {
    fn from(tensor: Tensor) -> Self {
        tensor.0
    }
}

/// Exact, shape-aware equality over the `f32` payload (Burn's `Tensor<B, D>`
/// deliberately has no `PartialEq`, so we define it on the host wrapper).
impl PartialEq for Tensor {
    fn eq(&self, other: &Self) -> bool {
        self.shape() == other.shape() && self.data() == other.data()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_rejects_shape_mismatch() {
        let err = Tensor::new([2, 3], vec![1.0, 2.0]).unwrap_err();
        assert!(matches!(err, Error::Shape { expected, actual }
            if expected == vec![2, 3] && actual == vec![2]));
    }

    #[test]
    fn new_rejects_overflowing_shape() {
        let err = Tensor::new([usize::MAX, 2], vec![0.0]).unwrap_err();
        assert!(matches!(err, Error::Validation(_)));
    }

    #[test]
    fn zeros_has_expected_len() {
        let t = Tensor::zeros([2, 3]).unwrap();
        assert_eq!(t.shape(), &[2, 3]);
        assert_eq!(t.len(), 6);
        assert!(t.data().iter().all(|&x| x == 0.0));
        assert!(matches!(
            Tensor::zeros([usize::MAX, 2]),
            Err(Error::Validation(_))
        ));
    }

    #[test]
    fn finiteness_is_observable() {
        assert!(Tensor::vector([1.0, -2.0]).is_finite());
        assert!(!Tensor::vector([1.0, f32::NAN]).is_finite());
        assert!(!Tensor::vector([f32::INFINITY]).is_finite());
    }

    #[test]
    fn round_trips_through_tensor_data() {
        let t = Tensor::new([2, 2], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let back = Tensor::from_data(t.clone().into_data()).unwrap();
        assert_eq!(t, back);
    }

    #[test]
    fn from_data_rejects_non_f32() {
        let i32_data = TensorData::new(vec![1i32, 2, 3], [3]);
        assert!(matches!(
            Tensor::from_data(i32_data),
            Err(Error::Validation(_))
        ));
    }

    #[test]
    fn try_from_rejects_length_mismatch() {
        // `from_bytes_vec` skips burn's own length assertion: four bytes is
        // one f32, but the shape declares five.
        let ragged = TensorData::from_bytes_vec(vec![0u8; 4], [5], burn_tensor::DType::F32);
        assert!(matches!(Tensor::try_from(ragged), Err(Error::Shape { .. })));
    }

    #[test]
    fn deserialize_rejects_non_f32() {
        let json = serde_json::to_string(&TensorData::new(vec![1i32, 2, 3], [3])).unwrap();
        assert!(serde_json::from_str::<Tensor>(&json).is_err());
    }

    #[test]
    fn serde_round_trip() {
        let t = Tensor::vector([1.0, -2.0, 3.5]);
        let json = serde_json::to_string(&t).unwrap();
        let back: Tensor = serde_json::from_str(&json).unwrap();
        assert_eq!(t, back);
    }
}