burn_pack/
tensor.rs

1//! Tensor-library-agnostic tensor entry for the burnpack format.
2//!
3//! The burnpack reader and writer operate on [`Tensor`] values, which carry only the
4//! format-level information (name, dtype, shape, optional param id) plus the raw
5//! little-endian [`Bytes`]. Keeping the bytes as [`Bytes`] (rather than a custom buffer
6//! type) integrates with the rest of the Burn ecosystem: a reader can hand out
7//! file-backed bytes ([`Bytes::from_file`]) for fast, lazy file-to-GPU transfers, while a
8//! writer simply consumes already-materialized bytes.
9
10use alloc::string::String;
11
12use burn_std::{Bytes, DType, Shape};
13
14/// A single tensor in a burnpack container, decoupled from any tensor library.
15///
16/// The [`bytes`](Self::bytes) field holds the tensor's data in little-endian layout,
17/// matching the element count implied by [`shape`](Self::shape) and [`dtype`](Self::dtype).
18/// When produced by a [`Reader`](crate::Reader) loading from a file, the bytes are
19/// file-backed and only read from disk when accessed.
20#[derive(Clone)]
21pub struct Tensor {
22    /// Fully-qualified tensor name (e.g. `"encoder.layer1.weight"`).
23    pub name: String,
24    /// Data type of the tensor.
25    pub dtype: DType,
26    /// Tensor shape.
27    pub shape: Shape,
28    /// Optional parameter id, used to preserve identities for stateful training.
29    pub param_id: Option<u64>,
30    /// The tensor's raw little-endian bytes.
31    pub bytes: Bytes,
32}
33
34impl Tensor {
35    /// Create a tensor entry from its metadata and raw bytes.
36    pub fn new(
37        name: String,
38        dtype: DType,
39        shape: impl Into<Shape>,
40        param_id: Option<u64>,
41        bytes: Bytes,
42    ) -> Self {
43        Self {
44            name,
45            dtype,
46            shape: shape.into(),
47            param_id,
48            bytes,
49        }
50    }
51
52    /// Number of raw bytes the tensor occupies.
53    pub fn byte_len(&self) -> usize {
54        self.bytes.len()
55    }
56}
burn_pack/tensor.rs

burn_pack/
tensor.rs
