acyclib 0.3.0

ML library for directed acyclic tensor graphs.
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use crate::device::{
    Device, OperationError,
    function::DeviceOperation,
    operation::{DiffableFromOutput, SparseAffineOps},
    tensor::{Shape, TensorRef},
};

#[derive(Clone)]
pub struct SparseAffineActivate<D: Device> {
    pub weights: TensorRef<D>,
    pub weights_shape: Shape,
    pub biases: Option<TensorRef<D>>,
    pub input_shape: Shape,
    pub indices: TensorRef<D>,
    pub values: Option<TensorRef<D>>,
    pub activation: DiffableFromOutput,
    pub output: TensorRef<D>,
}

impl<D: SparseAffineOps> DeviceOperation<D> for SparseAffineActivate<D> {
    fn opname(&self) -> String {
        "SparseAffineActivate".to_string()
    }

    fn execute(&self) -> Result<(), OperationError<D::DeviceError>> {
        let SparseAffineActivate { weights, weights_shape, biases, input_shape, indices, values, activation, output } =
            self;

        let weights = weights.dense();
        let indices = indices.sparse();
        let mut output = output.dense_mut();

        let biases = biases.as_ref().map(|b| b.dense());
        let biases = biases.as_deref();
        let values = values.as_ref().map(|v| v.dense());
        let values = values.as_deref();

        let batch_size = indices.batch_size();

        if batch_size != output.batch_size() {
            return Err(OperationError::MismatchedBatchSizes);
        }

        D::sparse_affine_activate(
            batch_size.unwrap_or(1),
            *activation,
            &weights.buf,
            *weights_shape,
            &indices.buf,
            values.map(|v| &v.buf),
            *input_shape,
            indices.nnz,
            biases.map(|b| &b.buf),
            biases.map(|b| b.batch_size.is_some()).unwrap_or(false),
            &mut output.buf,
        )
    }
}

#[derive(Clone)]
pub struct BackpropSparseAffineActivate<D: Device> {
    pub weights_grads: TensorRef<D>,
    pub weights_shape: Shape,
    pub biases_grads: Option<TensorRef<D>>,
    pub input_shape: Shape,
    pub indices: TensorRef<D>,
    pub values: Option<TensorRef<D>>,
    pub activation: DiffableFromOutput,
    pub output: TensorRef<D>,
    pub output_grads: TensorRef<D>,
}

impl<D: SparseAffineOps> DeviceOperation<D> for BackpropSparseAffineActivate<D> {
    fn opname(&self) -> String {
        "BackpropSparseAffineActivate".to_string()
    }

    fn execute(&self) -> Result<(), OperationError<<D as Device>::DeviceError>> {
        let BackpropSparseAffineActivate {
            weights_grads,
            weights_shape,
            biases_grads,
            input_shape,
            indices,
            values,
            activation,
            output,
            output_grads,
        } = self;

        let mut weights_grads = weights_grads.dense_mut();
        let indices = indices.sparse();
        let output = output.dense();
        let output_grads = output_grads.dense();

        let mut biases_grads = biases_grads.as_ref().map(|b| b.dense_mut());
        let biases_grads = biases_grads.as_deref_mut();
        let values = values.as_ref().map(|v| v.dense_mut());
        let values = values.as_deref();

        let biases_batched = biases_grads.as_ref().map(|b| b.batch_size.is_some()).unwrap_or(false);

        let batch_size = indices.batch_size();

        if batch_size != output.batch_size() || batch_size != output_grads.batch_size() {
            return Err(OperationError::MismatchedBatchSizes);
        }

        D::backprop_sparse_affine_activate(
            batch_size.unwrap_or(1),
            *activation,
            &mut weights_grads.buf,
            *weights_shape,
            &indices.buf,
            values.map(|v| &v.buf),
            *input_shape,
            indices.nnz,
            biases_grads.map(|b| &mut b.buf),
            biases_batched,
            &output.buf,
            &output_grads.buf,
        )
    }
}