forge-ml 0.1.0

A WebGPU-native machine learning runtime in Rust: train and run GPT-2 on any GPU wgpu reaches, with no CUDA and no Python
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

use std::fmt;

/// Row-major, contiguous tensor shape.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Shape(pub Vec<usize>);

impl Shape {
    pub fn new(dims: &[usize]) -> Self {
        Shape(dims.to_vec())
    }

    pub fn dims(&self) -> &[usize] {
        &self.0
    }

    pub fn rank(&self) -> usize {
        self.0.len()
    }

    /// Total number of elements.
    pub fn numel(&self) -> usize {
        self.0.iter().product()
    }

    pub fn dim(&self, i: usize) -> usize {
        self.0[i]
    }

    /// Row-major strides in elements.
    pub fn strides(&self) -> Vec<usize> {
        let mut strides = vec![1; self.0.len()];
        for i in (0..self.0.len().saturating_sub(1)).rev() {
            strides[i] = strides[i + 1] * self.0[i + 1];
        }
        strides
    }
}

impl fmt::Display for Shape {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self.0)
    }
}

impl From<&[usize]> for Shape {
    fn from(d: &[usize]) -> Self {
        Shape(d.to_vec())
    }
}

impl From<Vec<usize>> for Shape {
    fn from(d: Vec<usize>) -> Self {
        Shape(d)
    }
}

impl<const N: usize> From<[usize; N]> for Shape {
    fn from(d: [usize; N]) -> Self {
        Shape(d.to_vec())
    }
}