pub fn stack_to_tensor<const R: usize, const BR: usize, T, B>(
items: &[T],
device: &<B as BackendTypes>::Device,
) -> Tensor<B, BR>where
T: TensorConvertible<R, B>,
B: Backend,Expand description
Stages a whole batch of rows into one host buffer and uploads it as a single tensor.
Each item’s write_host_row payload is concatenated into one contiguous
Vec<f32>, which is then uploaded with a single Tensor::from_data
call. This is materially cheaper than converting each item to its own tensor
and calling Tensor::stack, which incurs one host→device upload per item
plus a concatenation kernel. Because both this function and the derived
TensorConvertible::to_tensor draw from the same
write_host_row/row_shape primitives, the batched layout is guaranteed
to match stack-ing the individual rows.
The produced tensor has rank BR = R + 1 and shape [items.len(), ..row],
i.e. a leading batch axis followed by the per-item row_shape.
§Type Parameters
R: rank of a single row.BR: rank of the batched tensor; must equalR + 1.T: the row type,TensorConvertible<R, B>.B: Burn backend.
§The BR = R + 1 contract
Stable Rust cannot express R + 1 in a const-generic position, so BR is a
separate parameter checked at runtime. This function is the single
chokepoint for that invariant: the leading assert_eq! runs before the
shape array is assembled, which is what makes the subsequent
shape[1..].copy_from_slice(&row) sound (it would panic on a length
mismatch otherwise).
§Panics
Panics if BR != R + 1.
§Examples
use burn::backend::Flex;
use burn::tensor::Tensor;
use rlevo_core::base::{stack_to_tensor, TensorConversionError, TensorConvertible};
#[derive(Clone)]
struct Point {
x: f32,
y: f32,
}
impl<B: burn::tensor::backend::Backend> TensorConvertible<1, B> for Point {
fn row_shape() -> [usize; 1] {
[2]
}
fn write_host_row(&self, buf: &mut Vec<f32>) {
buf.push(self.x);
buf.push(self.y);
}
fn from_tensor(_tensor: Tensor<B, 1>) -> Result<Self, TensorConversionError> {
unimplemented!()
}
}
type B = Flex;
let device = Default::default();
let items: Vec<Point> = vec![Point { x: 1.0, y: 2.0 }, Point { x: 3.0, y: 4.0 }];
let batched: Tensor<B, 2> = stack_to_tensor::<1, 2, Point, B>(&items, &device);
assert_eq!(batched.dims(), [2, 2]);