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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

//! Utility layer to give a tensor another shape.
//!
//! This layer should be used as in-place operation,
//! so the tensor that should be reshaped should be specified
//! as both input and output.
//!
//! Reshaping a tensor is required so that it becomes
//! usable for Layers that interpret meaning into the shape of
//! the tensor.
//!
//! A lot of layers interpret the last dimensions as NCHW,
//! where the letters stand for:
//!
//! - `N` : number of batch samples
//! - `C` : number of feature maps
//! - `H` : height
//! - `W` : width

use crate::capnp_util::*;
use crate::co::{IBackend, SharedTensor};
use crate::juice_capnp::reshape_config as capnp_config;
use crate::layer::*;
use crate::util::ArcLock;

#[derive(Debug, Clone)]
/// Reshape Utility Layer
pub struct Reshape {
    shape: Vec<usize>,
}

impl Reshape {
    /// Create a Reshape layer from a ReshapeConfig.
    pub fn from_config(config: &ReshapeConfig) -> Reshape {
        Reshape {
            shape: config.shape.clone(),
        }
    }
}

impl<B: IBackend> ILayer<B> for Reshape {
    fn compute_in_place(&self) -> bool {
        true
    }

    fn auto_output_blobs(&self) -> bool {
        false
    }

    fn reshape(
        &mut self,
        backend: ::std::rc::Rc<B>,
        input_data: &mut Vec<ArcLock<SharedTensor<f32>>>,
        input_gradient: &mut Vec<ArcLock<SharedTensor<f32>>>,
        weights_data: &mut Vec<ArcLock<SharedTensor<f32>>>,
        weights_gradient: &mut Vec<ArcLock<SharedTensor<f32>>>,
        output_data: &mut Vec<ArcLock<SharedTensor<f32>>>,
        output_gradient: &mut Vec<ArcLock<SharedTensor<f32>>>,
    ) {
        output_data[0].write().unwrap().resize(&self.shape).unwrap();
        output_gradient[0].write().unwrap().resize(&self.shape).unwrap();
    }
}

impl<B: IBackend> ComputeOutput<f32, B> for Reshape {
    fn compute_output(
        &self,
        backend: &B,
        _weights: &[&SharedTensor<f32>],
        input_data: &[&SharedTensor<f32>],
        output_data: &mut [&mut SharedTensor<f32>],
    ) {
    }
}

impl<B: IBackend> ComputeInputGradient<f32, B> for Reshape {
    fn compute_input_gradient(
        &self,
        backend: &B,
        weights_data: &[&SharedTensor<f32>],
        output_data: &[&SharedTensor<f32>],
        output_gradients: &[&SharedTensor<f32>],
        input_data: &[&SharedTensor<f32>],
        input_gradients: &mut [&mut SharedTensor<f32>],
    ) {
    }
}

impl<B: IBackend> ComputeParametersGradient<f32, B> for Reshape {}

#[derive(Debug, Clone)]
/// Specifies configuration parameters for a Reshape Layer.
pub struct ReshapeConfig {
    /// The target shape that the input should assume.
    ///
    /// Preceding dimensions are treated as independent inputs
    ///
    /// Defaults to `1`
    pub shape: Vec<usize>,
}

impl ReshapeConfig {
    /// Create a ReshapeConfig that describes a Reshape layer with a provided shape.
    pub fn of_shape(shape: &[usize]) -> ReshapeConfig {
        ReshapeConfig {
            shape: shape.to_owned(),
        }
    }
}

impl<'a> CapnpWrite<'a> for ReshapeConfig {
    type Builder = capnp_config::Builder<'a>;

    /// Write the ReshapeConfig into a capnp message.
    fn write_capnp(&self, builder: &mut Self::Builder) {
        let mut shape = builder.reborrow().init_shape(self.shape.len() as u32);
        for (i, dim) in self.shape.iter().enumerate() {
            shape.set(i as u32, *dim as u64);
        }
    }
}

impl<'a> CapnpRead<'a> for ReshapeConfig {
    type Reader = capnp_config::Reader<'a>;

    fn read_capnp(reader: Self::Reader) -> Self {
        let read_shape = reader.get_shape().unwrap();
        let mut shape = Vec::new();
        for i in 0..read_shape.len() {
            shape.push(read_shape.get(i) as usize)
        }

        ReshapeConfig { shape: shape }
    }
}

impl Into<LayerType> for ReshapeConfig {
    fn into(self) -> LayerType {
        LayerType::Reshape(self)
    }
}