fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
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

//! Derivative of row selection.

use crate::autograd::Backward;
use crate::cuda::kernels;
use crate::tensor::Tensor;
use crate::tensor::storage::Storage;

/// Scatter-add each row's gradient into its source row, so repeats accumulate.
pub struct IndexSelectBackward {
    pub ids: Vec<usize>,
    pub rows: usize,
    pub cols: usize,
}

impl Backward for IndexSelectBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        if let Storage::Cuda(g) = grad.storage() {
            let ids_i32: Vec<i32> = self.ids.iter().map(|&i| i as i32).collect();
            let ids_buf = kernels::upload_ids(&ids_i32).expect("cuda upload ids");
            let table =
                kernels::embedding_backward(&ids_buf, g, self.ids.len(), self.cols, self.rows)
                    .expect("cuda embedding backward");
            return vec![Tensor::raw(
                Storage::Cuda(table),
                vec![self.rows, self.cols],
                grad.device(),
            )];
        }

        let g = grad.to_vec();
        let mut table = vec![0.0f32; self.rows * self.cols];
        for (position, &id) in self.ids.iter().enumerate() {
            let (dst, src) = (id * self.cols, position * self.cols);
            for c in 0..self.cols {
                table[dst + c] += g[src + c];
            }
        }
        vec![Tensor::from_vec(table, &[self.rows, self.cols]).to(grad.device())]
    }
    fn name(&self) -> &'static str {
        "IndexSelect"
    }
}