Skip to main content

stack_to_tensor

Function stack_to_tensor 

Source
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 equal R + 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]);