1use proc_macro::TokenStream;
2use quote::quote;
3
4fn get_layers_and_ident(input: syn::DeriveInput) ->
5 (syn::Type, syn::Ident)
6{
7 let ident = input.ident;
8 if let syn::Data::Struct(data) = input.data {
9 if let syn::Fields::Named(fields) = data.fields {
10 if fields.named.len() != 1 {
11 panic!("Could not construct neural \
12 network: {} struct has too \
13 many fields!", ident);
14 }
15
16 let layers_field = fields.named
17 .into_iter().next().unwrap();
18 let field_ident = layers_field.ident;
19 if field_ident.is_none() {
20 panic!("Could not construct neural network: \
21 {} struct does not have `layers` \
22 field!", ident);
23 }
24 let fieldname = field_ident
25 .unwrap().to_string();
26
27 if String::from("layers") != fieldname {
28 panic!("Could not construct neural network: \
29 {} struct does not have `layers` \
30 field!", ident);
31 }
32
33 return (layers_field.ty, ident);
34 } else {
35 panic!("Could not construct neural network: \
36 {} struct does not have `layers` \
37 field!", ident);
38 }
39 } else {
40 panic!("Could not construct neural network: \
41 {} is not a struct!", ident);
42 }
43}
44
45#[proc_macro_attribute]
46pub fn neural_network(_attr: TokenStream, item: TokenStream) -> TokenStream {
47 let input: syn::DeriveInput =
48 syn::parse(item.clone()).unwrap();
49 let (ty, ident) = get_layers_and_ident(input);
50
51 let layers = if let syn::Type::Tuple(layers) = ty {
52 layers.elems
53 } else {
54 panic!("Could not construct neural network: \
55 `layers` is not a tuple!");
56 };
57
58 let mut layer_idents: Vec<syn::Path> = Vec::new();
59 for layer in layers.into_iter() {
60 layer_idents.push(match layer {
61 syn::Type::Path(path) => path.path,
62 _ => panic!("Could not construct neural network: \
63 Invalid layer type!"),
64 });
65 }
66
67 let layer_idents: Vec<syn::Path> = layer_idents
68 .into_iter().collect();
69
70 let first_layer = &layer_idents[0];
71 let last_layer = &layer_idents[
72 layer_idents.len()-1];
73
74 let mut new_list = proc_macro2::TokenStream::new();
75 for layer_ident in layer_idents.iter() {
76 new_list.extend(quote!{
77 #layer_ident::new(),
78 });
79 }
80
81 let mut layers_string = String::from("[");
82 for i in 0..layer_idents.len() {
83 let layer_ident = &layer_idents[i];
84 layers_string += "\"";
85 layers_string += "e!(#layer_ident).to_string();
86 layers_string += "\"";
87 if i < layer_idents.len()-1 {
88 layers_string += ", ";
89 }
90 }
91 layers_string += "]";
92
93 let mut params_cnt_sum = proc_macro2::TokenStream::new();
94 for i in 0..layer_idents.len() {
95 let layer_ident = &layer_idents[i];
96 params_cnt_sum.extend(quote!{#layer_ident::PARAMS_CNT});
97 if i < layer_idents.len()-1 {
98 params_cnt_sum.extend(quote!{ + });
99 }
100 }
101
102 let mut d_offsets: Vec<proc_macro2::TokenStream> =
103 Vec::with_capacity(layer_idents.len()+1);
104 d_offsets.push(quote!{0});
105 for i in 0..layer_idents.len() {
106 let layer_ident = &layer_idents[i];
107 d_offsets.push(d_offsets[i].clone());
108 d_offsets[i+1].extend(quote!{
109 + #layer_ident::PARAMS_CNT
110 });
111 }
112
113 let mut eval_all_layers = proc_macro2::TokenStream::new();
114 for i in 0..layer_idents.len() {
115 let layer_ident = &layer_idents[i];
116
117 let old_offset = &d_offsets[i];
118 let offset = &d_offsets[i+1];
119 eval_all_layers.extend(quote!{
120 let x = unsafe {
121 #layer_ident::eval_unchecked(
122 &p[#old_offset..#offset], x)
123 };
124 });
125 }
126
127 let mut forward_all_layers = proc_macro2::TokenStream::new();
128 let mut backward_all_layers = proc_macro2::TokenStream::new();
129 for i in 0..layer_idents.len() {
130 let old_offset = &d_offsets[i];
131 let offset = &d_offsets[i+1];
132 let idx: syn::Index = i.into();
133
134 forward_all_layers.extend(quote!{
135 let x = self.layers.#idx.forward(
136 &p[#old_offset..#offset], x);
137 });
138 backward_all_layers.extend(quote!{
139 self.layers.#idx.backward(
140 &p[#old_offset..#offset]);
141 });
142 }
143
144 let mut compute_jacobian = proc_macro2::TokenStream::new();
145 compute_jacobian.extend(quote!{
146 let mut jm: DMatrix<f64> =
147 DMatrix::from_element_generic(
148 nalgebra::base::dimension::Dyn(Self::NEURONS_OUT),
149 nalgebra::base::dimension::Dyn(Self::PARAMS_CNT), 0f64);
150 let m: DMatrix<f64> = {
151 let mut m: DMatrix<f64> =
152 DMatrix::from_element_generic(
153 nalgebra::base::dimension::Dyn(Self::NEURONS_OUT),
154 nalgebra::base::dimension::Dyn(Self::NEURONS_OUT), 0f64);
155 m.fill_diagonal(1f64);
156
157 m
158 };
159 let mut offset: usize = Self::PARAMS_CNT;
160 });
161
162 for i in (1..layer_idents.len()).rev() {
163 let layer_ident = &layer_idents[i];
164 let idx: syn::Index = i.into();
165 let prev_idx: syn::Index = (i-1).into();
166 compute_jacobian.extend(quote!{
167 let jf = &m * self.layers.#idx.chain_end(
168 &self.layers.#prev_idx.signal);
169 offset -= #layer_ident::PARAMS_CNT;
170 for i in offset..offset+#layer_ident::PARAMS_CNT {
171 jm.set_column(i, &jf.index((.., i - offset)));
172 }
173
174 let m = m * self.layers.#idx.chain_element();
175 });
176 }
177 let idx: syn::Index = 0.into();
178 compute_jacobian.extend(quote!{
179 let jf = m * self.layers.#idx.chain_end(&x);
180 for i in 0..#first_layer::PARAMS_CNT {
181 jm.set_column(i, &jf.index((.., i)));
182 }
183 });
184
185 let mut extend_by_initial_params = proc_macro2::TokenStream::new();
186 for i in 0..layer_idents.len() {
187 let layer_ident = &layer_idents[i];
188
189 extend_by_initial_params.extend(quote!{
190 p.append(&mut #layer_ident::default_initial_params());
191 });
192 }
193
194 let network_trait_impl = quote! {
195 impl Network for #ident {
196 const PARAMS_CNT: usize = #params_cnt_sum;
197 const NEURONS_IN: usize = #first_layer::NEURONS_IN;
198 const NEURONS_OUT: usize = #last_layer::NEURONS_OUT;
199
200 fn new() -> Self {
201 Self {
202 layers: (#new_list),
203 }
204 }
205
206 fn layers_info() -> &'static str {
207 #layers_string
208 }
209
210 fn eval(p: &[f64], x: DVector<f64>) ->
211 DVector<f64>
212 {
213 assert_eq!(p.len(), Self::PARAMS_CNT);
214 assert_eq!(x.len(), Self::NEURONS_IN);
215
216 #eval_all_layers
217
218 x
219 }
220
221 fn forward(&mut self, p: &[f64], x: DVector<f64>) ->
222 DVector<f64>
223 {
224 assert_eq!(p.len(), Self::PARAMS_CNT);
225 assert_eq!(x.len(), Self::NEURONS_IN);
226
227 #forward_all_layers
228
229 x
230 }
231
232 fn backward(&mut self, p: &[f64])
233 {
234 assert_eq!(p.len(), Self::PARAMS_CNT);
235
236 #backward_all_layers
237 }
238
239 fn jacobian(&mut self, x: &DVector<f64>) ->
240 DMatrix<f64>
241 {
242 assert_eq!(x.len(), Self::NEURONS_IN);
243
244 #compute_jacobian
245
246 jm
247 }
248
249 fn default_initial_params() -> Vec<f64> {
250 let mut p: Vec<f64> =
251 Vec::with_capacity(Self::PARAMS_CNT);
252
253 #extend_by_initial_params
254
255 p
256 }
257 }
258 };
259
260 let mut output =
261 proc_macro2::TokenStream::from(item);
262 output.extend(network_trait_impl);
263
264 proc_macro::TokenStream::from(output)
265}