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
|
#include "utils.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
RNG rng_init(uint64_t seed) {
RNG r;
r.state = seed ? seed : 1;
return r;
}
uint64_t rng_next(RNG *r) {
uint64_t x = r->state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
r->state = x;
return x;
}
double rng_uniform(RNG *r) {
return (double)(rng_next(r) >> 11) / (double)(UINT64_C(1) << 53);
}
double rng_normal(RNG *r) {
double u1 = rng_uniform(r);
double u2 = rng_uniform(r);
if (u1 < 1e-300) u1 = 1e-300;
return sqrt(-2.0 * log(u1)) * cos(2.0 * M_PI * u2);
}
void *xmalloc(size_t n) {
void *p = malloc(n);
if (!p) { fprintf(stderr, "fatal: malloc(%zu) failed\n", n); abort(); }
return p;
}
void *xcalloc(size_t count, size_t size) {
void *p = calloc(count, size);
if (!p) { fprintf(stderr, "fatal: calloc(%zu, %zu) failed\n", count, size); abort(); }
return p;
}
double *vec_alloc(int n) { return (double *)xmalloc((size_t)n * sizeof(double)); }
double *vec_zeros(int n) { return (double *)xcalloc((size_t)n, sizeof(double)); }
void vec_zero(double *v, int n) { memset(v, 0, (size_t)n * sizeof(double)); }
void vec_copy(const double *src, double *dst, int n) { memcpy(dst, src, (size_t)n * sizeof(double)); }
void vec_add_scaled(double *dst, double alpha, const double *v, int n) {
for (int i = 0; i < n; i++) dst[i] += alpha * v[i];
}
double vec_dot(const double *a, const double *b, int n) {
double s = 0.0;
for (int i = 0; i < n; i++) s += a[i] * b[i];
return s;
}
void mat_vec(const double *M, const double *x, double *dst, int rows, int cols) {
for (int i = 0; i < rows; i++) {
double s = 0.0;
for (int j = 0; j < cols; j++) s += M[i * cols + j] * x[j];
dst[i] = s;
}
}
void mat_vec_T(const double *M, const double *v, double *dst, int rows, int cols) {
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++)
dst[j] += M[i * cols + j] * v[i];
}
void mat_outer_add(double *M, double alpha,
const double *a, const double *b, int rows, int cols) {
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++)
M[i * cols + j] += alpha * a[i] * b[j];
}
Workspace workspace_alloc(int D, int H, int nparams) {
Workspace ws;
ws.x = vec_alloc(D + 1);
ws.h_pre = vec_alloc(H);
ws.h = vec_alloc(H);
ws.dh = vec_alloc(H);
ws.dh_pre = vec_alloc(H);
ws.dx = vec_alloc(D + 1);
ws.neg_a = vec_alloc(D);
ws.vjp_z = vec_alloc(D);
ws.vjp_theta = vec_alloc(nparams);
return ws;
}
void workspace_free(Workspace *ws) {
free(ws->x);
free(ws->h_pre);
free(ws->h);
free(ws->dh);
free(ws->dh_pre);
free(ws->dx);
free(ws->neg_a);
free(ws->vjp_z);
free(ws->vjp_theta);
}
|