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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

//! # Shape Helpers

use burn::{
    prelude::Shape,
    tensor::{
        AsIndex,
        wrap_index,
    },
};

/// Compute the ravel index for the given coordinates.
///
/// This returns the row-major order raveling.
///
/// # Arguments
/// - `coords`: must be the same size as `self.rank()`.
///
/// # Returns
/// - the ravel offset index.
pub fn ravel_shape<I: AsIndex>(
    shape: &Shape,
    coords: &[I],
) -> usize {
    ravel_dims(shape.as_slice(), coords)
}

/// Compute the ravel index for the given coordinates.
///
/// This returns the row-major order raveling.
///
/// # Arguments
/// - `dims`: the dimensions of the shape.
/// - `coords`: must be the same size as `self.rank()`.
///
/// # Returns
/// - the ravel offset index.
pub fn ravel_dims<I: AsIndex>(
    dims: &[usize],
    coords: &[I],
) -> usize {
    assert_eq!(
        dims.len(),
        coords.len(),
        "Shape rank mismatch: expected {}, got {}",
        dims.len(),
        coords.len(),
    );

    let mut ravel_idx = 0;
    let mut stride = 1;

    for i in (0..dims.len()).rev() {
        let dim = dims[i];
        let coord = wrap_index(coords[i], dim);

        ravel_idx += coord * stride;
        stride *= dim;
    }

    ravel_idx
}