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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

//! Element-wise maths on one tensor.

use crate::autograd::ops::unary::{
    AbsBackward, ClampBackward, ExpBackward, LogBackward, NegBackward, PowBackward, SqrtBackward,
};
use crate::cuda::{ffi, kernels};
use crate::tensor::Tensor;
use crate::tensor::storage::Storage;

use super::{scalar_op, unary_op};

impl Tensor {
    /// `-self`.
    pub fn neg(&self) -> Tensor {
        unary_op(self, |x| -x, ffi::fastnn_cuda_neg).with_grad(&[self], || NegBackward)
    }

    /// `e^self`.
    pub fn exp(&self) -> Tensor {
        let out = unary_op(self, f32::exp, ffi::fastnn_cuda_exp);
        let saved = out.clone();
        out.with_grad(&[self], || ExpBackward { output: saved })
    }

    /// Natural log. Undefined for non-positive inputs.
    pub fn log(&self) -> Tensor {
        unary_op(self, f32::ln, ffi::fastnn_cuda_log).with_grad(&[self], || LogBackward {
            input: self.detach(),
        })
    }

    /// `√self`.
    pub fn sqrt(&self) -> Tensor {
        let out = unary_op(self, f32::sqrt, ffi::fastnn_cuda_sqrt);
        let saved = out.clone();
        out.with_grad(&[self], || SqrtBackward { output: saved })
    }

    /// `|self|`. The derivative at exactly 0 is taken as 0.
    pub fn abs(&self) -> Tensor {
        unary_op(self, f32::abs, ffi::fastnn_cuda_abs).with_grad(&[self], || AbsBackward {
            input: self.detach(),
        })
    }

    /// `self` raised to a constant power.
    pub fn powf(&self, exponent: f32) -> Tensor {
        scalar_op(
            self,
            exponent,
            |x, p| x.powf(p),
            ffi::fastnn_cuda_pow_scalar,
        )
        .with_grad(&[self], || PowBackward {
            input: self.detach(),
            exponent,
        })
    }

    /// `self * self`.
    pub fn square(&self) -> Tensor {
        self.mul(self)
    }

    /// Clip every element into `[lo, hi]`. The gradient is zero outside the range.
    pub fn clamp(&self, lo: f32, hi: f32) -> Tensor {
        let out = match self.storage() {
            Storage::Cpu(data) => {
                Tensor::from_vec(super::map(data, |x| x.clamp(lo, hi)), self.shape())
            }
            Storage::Cuda(buf) => {
                let out = kernels::clamp(buf, lo, hi, self.numel()).expect("cuda clamp");
                Tensor::raw(Storage::Cuda(out), self.shape().to_vec(), self.device())
            }
        };
        out.with_grad(&[self], || ClampBackward {
            input: self.detach(),
            lo,
            hi,
        })
    }
}