furnace/
tensor.rs

1//! Tensors as the server sees them: named inputs and outputs of a model, with a datatype and a
2//! shape, held in host memory between the wire formats and the engine that runs the model.
3//!
4//! Every request reaches a model as [`HostTensor`]s, one per input, batch first. The flat
5//! `Vec<f32>` entry points of `Model` (`predict*`) wrap their buffer as one such tensor without
6//! copying it.
7
8use serde::{Deserialize, Serialize};
9
10use crate::error::{ModelError, Result};
11
12/// Element types a model's inputs and outputs can have
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum DataType {
15    Fp32,
16    Int64,
17    Int32,
18}
19
20impl DataType {
21    /// The Open Inference Protocol name: `FP32`, `INT64`, `INT32`
22    pub fn v2_name(self) -> &'static str {
23        match self {
24            Self::Fp32 => "FP32",
25            Self::Int64 => "INT64",
26            Self::Int32 => "INT32",
27        }
28    }
29
30    /// The type an Open Inference Protocol name stands for; `None` for one the server does not
31    /// take
32    pub fn from_v2_name(name: &str) -> Option<Self> {
33        match name {
34            "FP32" => Some(Self::Fp32),
35            "INT64" => Some(Self::Int64),
36            "INT32" => Some(Self::Int32),
37            _ => None,
38        }
39    }
40
41    /// Bytes per element in the raw (binary) format
42    pub fn size(self) -> usize {
43        match self {
44            Self::Fp32 | Self::Int32 => 4,
45            Self::Int64 => 8,
46        }
47    }
48
49    /// Whether elements are integers (held as `i64` in [`Values::I64`])
50    pub fn is_int(self) -> bool {
51        matches!(self, Self::Int64 | Self::Int32)
52    }
53}
54
55/// A model input or output: its name, element type and per-item shape (without the batch
56/// dimension; `None` where the graph leaves a dimension open)
57#[derive(Debug, Clone, PartialEq)]
58pub struct IoSpec {
59    pub name: String,
60    pub dtype: DataType,
61    pub shape: Vec<Option<usize>>,
62}
63
64impl IoSpec {
65    /// A tensor of `dtype` with per-item `shape` (`None` for a dimension each request fixes)
66    pub fn new(name: &str, dtype: DataType, shape: &[Option<usize>]) -> Self {
67        Self {
68            name: name.to_string(),
69            dtype,
70            shape: shape.to_vec(),
71        }
72    }
73
74    /// An `f32` tensor with a fixed per-item shape
75    pub fn fp32(name: &str, shape: &[usize]) -> Self {
76        Self {
77            name: name.to_string(),
78            dtype: DataType::Fp32,
79            shape: shape.iter().map(|&d| Some(d)).collect(),
80        }
81    }
82
83    /// The per-item shape if every dimension is fixed
84    pub fn static_shape(&self) -> Option<Vec<usize>> {
85        self.shape.iter().copied().collect()
86    }
87
88    /// For metadata: the full shape with `-1` for the batch and for open dimensions
89    pub fn v2_shape(&self) -> Vec<i64> {
90        std::iter::once(-1)
91            .chain(self.shape.iter().map(|d| d.map_or(-1, |d| d as i64)))
92            .collect()
93    }
94}
95
96/// The legacy single-tensor description in `GET /model/info` (`input_spec`, `output_spec`):
97/// per-item shape, open dimensions as 0
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct TensorSpec {
100    pub shape: Vec<usize>,
101    /// `float32`, `int64` or `int32`
102    pub dtype: String,
103    pub min_value: Option<f32>,
104    pub max_value: Option<f32>,
105}
106
107/// Description of a model input or output in `GET /model/info`
108#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
109pub struct TensorInfo {
110    pub name: String,
111    /// `FP32`, `INT64` or `INT32`
112    pub datatype: String,
113    /// Full shape; `-1` for the batch dimension and for dimensions fixed per request
114    pub shape: Vec<i64>,
115}
116
117impl From<&IoSpec> for TensorInfo {
118    fn from(spec: &IoSpec) -> Self {
119        Self {
120            name: spec.name.clone(),
121            datatype: spec.dtype.v2_name().to_string(),
122            shape: spec.v2_shape(),
123        }
124    }
125}
126
127/// The number of elements of a tensor of `shape`; `None` when it does not fit in a `usize`,
128/// as a client's shape may not
129pub(crate) fn element_count(shape: &[usize]) -> Option<usize> {
130    shape.iter().try_fold(1usize, |n, &d| n.checked_mul(d))
131}
132
133/// Tensor elements in host memory
134#[derive(Debug, Clone, PartialEq)]
135pub enum Values {
136    F32(Vec<f32>),
137    I64(Vec<i64>),
138}
139
140impl Values {
141    /// Number of elements
142    pub fn len(&self) -> usize {
143        match self {
144            Self::F32(v) => v.len(),
145            Self::I64(v) => v.len(),
146        }
147    }
148
149    /// Whether there are no elements
150    pub fn is_empty(&self) -> bool {
151        self.len() == 0
152    }
153}
154
155/// A tensor in host memory: full shape (batch first) and row-major elements.
156///
157/// What [`TensorModel::run`](crate::TensorModel::run) takes and returns. furnace-burn turns one
158/// into a Burn tensor and back (its `HostTensorExt`).
159#[derive(Debug, Clone, PartialEq)]
160pub struct HostTensor {
161    pub shape: Vec<usize>,
162    pub values: Values,
163}
164
165impl HostTensor {
166    /// Items along the batch dimension; 0 for a tensor without one (an empty shape), which no
167    /// model accepts
168    pub fn batch(&self) -> usize {
169        self.shape.first().copied().unwrap_or(0)
170    }
171
172    /// Items `start..end` along the batch dimension
173    pub fn slice_batch(&self, start: usize, end: usize) -> Self {
174        let per_item: usize = self.shape[1..].iter().product();
175        let range = start * per_item..end * per_item;
176        let mut shape = self.shape.clone();
177        shape[0] = end - start;
178        Self {
179            shape,
180            values: match &self.values {
181                Values::F32(v) => Values::F32(v[range].to_vec()),
182                Values::I64(v) => Values::I64(v[range].to_vec()),
183            },
184        }
185    }
186
187    /// Join tensors of the same type and per-item shape along the batch dimension
188    pub fn concat_batch(parts: Vec<Self>) -> Result<Self> {
189        let mut parts = parts.into_iter();
190        let mut joined = parts
191            .next()
192            .ok_or_else(|| ModelError::InferenceFailed("nothing to join".to_string()))?;
193        for part in parts {
194            if part.shape[1..] != joined.shape[1..] {
195                return Err(ModelError::InferenceFailed(format!(
196                    "chunks of one batch returned shapes {:?} and {:?}",
197                    joined.shape, part.shape
198                ))
199                .into());
200            }
201            joined.shape[0] += part.shape[0];
202            match (&mut joined.values, part.values) {
203                (Values::F32(a), Values::F32(b)) => a.extend(b),
204                (Values::I64(a), Values::I64(b)) => a.extend(b),
205                _ => {
206                    return Err(ModelError::InferenceFailed(
207                        "chunks of one batch returned different types".to_string(),
208                    )
209                    .into());
210                }
211            }
212        }
213        Ok(joined)
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn batches_split_and_join() {
223        let tensor = HostTensor {
224            shape: vec![3, 2],
225            values: Values::I64(vec![1, 2, 3, 4, 5, 6]),
226        };
227        let parts = vec![tensor.slice_batch(0, 1), tensor.slice_batch(1, 3)];
228        assert_eq!(parts[1].shape, [2, 2]);
229        assert_eq!(HostTensor::concat_batch(parts).unwrap(), tensor);
230    }
231
232    #[test]
233    fn metadata_shapes() {
234        let spec = IoSpec {
235            name: "ids".to_string(),
236            dtype: DataType::Int64,
237            shape: vec![Some(256), None],
238        };
239        assert_eq!(spec.v2_shape(), [-1, 256, -1]);
240        assert_eq!(spec.static_shape(), None);
241        assert_eq!(TensorInfo::from(&spec).datatype, "INT64");
242        assert_eq!(DataType::from_v2_name("INT32"), Some(DataType::Int32));
243        assert_eq!(DataType::from_v2_name("FP16"), None);
244    }
245}
furnace/tensor.rs

furnace/
tensor.rs
