ts_wire/
tensor.rs

1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! Tensors as the Python codec carries them (`MsgpackEncoder` in
22//! `io_struct.py`).
23//!
24//! A tensor is the array `[dtype, shape, data]`. `dtype` is a torch dtype
25//! name (`"bfloat16"`) or a numpy type string (`"<f4"`); `data` is the
26//! contiguous little-endian bytes, either inline as an ext value with code
27//! [`RAW_VIEW_EXT`] or, from [`crate::msgpack::INLINE_THRESHOLD`] bytes up,
28//! the index of the out-of-band frame that holds them.
29
30use bytes::Bytes;
31
32use crate::{
33    codec::Wire,
34    error::{Error, Result},
35    msgpack::{Decoder, Encoder, Kind},
36};
37
38/// Ext code of an inline tensor buffer (`CUSTOM_TYPE_RAW_VIEW`). Codes 1 and
39/// 2 are reserved for pickled payloads, which this protocol never carries.
40pub const RAW_VIEW_EXT: i8 = 3;
41
42/// A CPU tensor's metadata and bytes.
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct Tensor {
45    /// torch dtype name or numpy type string.
46    pub dtype: String,
47    /// Dimensions, outermost first.
48    pub shape: Vec<i64>,
49    /// Contiguous element bytes. Shares the received frame; no copy.
50    pub data: Bytes,
51}
52
53impl Wire for Tensor {
54    fn encode(&self, encoder: &mut Encoder) {
55        encoder.array_len(3);
56        encoder.str(&self.dtype);
57        self.shape.encode(encoder);
58        encoder.buffer(RAW_VIEW_EXT, &self.data);
59    }
60
61    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
62        let len = decoder.array_len()?;
63        if len != 3 {
64            return Err(Error::Invalid(format!(
65                "a tensor is [dtype, shape, data], found {len} elements"
66            )));
67        }
68        let dtype = String::decode(decoder)?;
69        let shape = Vec::decode(decoder)?;
70        let data = match decoder.peek_kind()? {
71            Kind::Ext => {
72                let (code, data) = decoder.ext()?;
73                if code != RAW_VIEW_EXT {
74                    return Err(Error::Invalid(format!(
75                        "tensor data has ext code {code}, expected {RAW_VIEW_EXT}"
76                    )));
77                }
78                data
79            }
80            Kind::Int => {
81                let index = decoder.integer()?;
82                decoder.frame(index)?
83            }
84            _ => {
85                return Err(Error::Invalid(
86                    "tensor data must be an inline ext buffer or a frame index".to_owned(),
87                ))
88            }
89        };
90        Ok(Self { dtype, shape, data })
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::codec::{decode_frames, encode_frames};
98
99    fn tensor(len: usize) -> Tensor {
100        Tensor {
101            dtype: "uint8".to_owned(),
102            shape: vec![len as i64],
103            data: Bytes::from(vec![7u8; len]),
104        }
105    }
106
107    #[test]
108    fn small_tensors_inline_and_large_ones_get_a_frame() {
109        let small = encode_frames(&tensor(255));
110        assert_eq!(small.len(), 1);
111        assert_eq!(decode_frames::<Tensor>(&small).unwrap(), tensor(255));
112
113        let large = encode_frames(&tensor(256));
114        assert_eq!(large.len(), 2);
115        assert_eq!(large[1].len(), 256);
116        assert_eq!(decode_frames::<Tensor>(&large).unwrap(), tensor(256));
117    }
118
119    #[test]
120    fn a_frame_index_past_the_message_fails() {
121        let mut frames = encode_frames(&tensor(300));
122        frames.truncate(1);
123        assert!(decode_frames::<Tensor>(&frames).is_err());
124    }
125}
ts_wire/tensor.rs

ts_wire/
tensor.rs
