bunsen 0.22.0

bunsen is a batteries included common library for burn
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

//! # `TensorData` View Wrappers

use std::ops::Index;

use burn::{
    prelude::TensorData,
    tensor::{
        AsIndex,
        Element,
    },
};

use crate::zspace::ravel_dims;

/// Ravel Index View for a `TensorData`.
#[derive(Debug)]
pub struct TensorDataIndexView<'a, E: Element> {
    data: &'a TensorData,
    _phantom: std::marker::PhantomData<&'a E>,
}

impl<'a, E: Element> TensorDataIndexView<'a, E> {
    /// Get an indexed view of the data.
    pub fn view(data: &'a TensorData) -> TensorDataIndexView<'a, E> {
        TensorDataIndexView {
            data,
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexView<'a, E> {
    type Output = E;

    fn index(
        &self,
        index: &[I],
    ) -> &Self::Output {
        &self.data.as_slice::<E>().unwrap()[ravel_dims(&self.data.shape, index)]
    }
}