burn-rmexp-dyntensor 0.0.5

burn remote experiment dyntensor
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

//! # Tensor Operations
use burn::Tensor;
use burn::prelude::Backend;
use burn::tensor::{BasicOps, Slice};

/// Provides a dynamic version of [`Tensor::slice`].
pub fn slice_dyn<B: Backend, const R: usize, K: BasicOps<B>>(
    tensor: Tensor<B, R, K>,
    slices: &[Slice],
) -> Tensor<B, R, K> {
    let mut tensor = tensor;
    for (dim, slice) in slices.iter().enumerate() {
        tensor = tensor.slice_dim(dim, *slice);
    }
    tensor
}