burn_rl/transition_buffer/
slice_access.rs

1use burn_core::prelude::*;
2
3/// Trait for types that support tensor-like slice operations,
4/// enabling storage in a [`TransitionBuffer`](super::TransitionBuffer).
5///
6/// Implement this trait for any type that wraps tensors and can be stored
7/// in a replay buffer. The buffer uses these operations for:
8/// - Pre-allocating storage (`zeros_like`)
9/// - Writing transitions (`slice_assign_inplace`)
10/// - Sampling batches (`select`)
11pub trait SliceAccess<B: Backend>: Clone + Sized {
12    /// Create zeroed storage matching the shape of `sample` but with `capacity` rows
13    /// along the first dimension.
14    fn zeros_like(sample: &Self, capacity: usize, device: &B::Device) -> Self;
15
16    /// Select rows at the given indices along the specified dimension.
17    fn select(self, dim: usize, indices: Tensor<B, 1, Int>) -> Self;
18
19    /// Assign `value` at row `index` along the first dimension, in place.
20    fn slice_assign_inplace(&mut self, index: usize, value: Self);
21}
22
23impl<B: Backend> SliceAccess<B> for Tensor<B, 2> {
24    fn zeros_like(sample: &Self, capacity: usize, device: &B::Device) -> Self {
25        let feature_dim = sample.dims()[1];
26        Tensor::zeros([capacity, feature_dim], device)
27    }
28
29    fn select(self, dim: usize, indices: Tensor<B, 1, Int>) -> Self {
30        Tensor::select(self, dim, indices)
31    }
32
33    fn slice_assign_inplace(&mut self, index: usize, value: Self) {
34        self.inplace(|t| t.slice_assign(index..index + 1, value));
35    }
36}
burn_rl/transition_buffer/slice_access.rs

burn_rl/transition_buffer/
slice_access.rs
