1 files changed, 94 insertions, 0 deletions
diff --git a/src/dynmlp.c b/src/dynmlp.c
new file mode 100644
index 0000000..70c9b75
--- /dev/null
+++ b/src/dynmlp.c
@@ -0,0 +1,94 @@
+#include "dynmlp.h"
+
+#include <math.h>
+
+static void xavier_init(double *w, int fan_in, int fan_out, RNG *r) {
+    double limit = sqrt(6.0 / (fan_in + fan_out));
+    int n = fan_in * fan_out;
+    for (int i = 0; i < n; i++)
+        w[i] = (2.0 * rng_uniform(r) - 1.0) * limit;
+}
+
+int dynmlp_nparams(int D, int H) {
+    return (D + 1) * H + H + H * D + D;
+}
+
+void dynmlp_init(DynMLP *net, int D, int H, double *theta, RNG *r) {
+    net->D = D;
+    net->H = H;
+    net->nparams = dynmlp_nparams(D, H);
+    xavier_init(theta + DYNMLP_W1(D, H), D + 1, H, r);
+    vec_zero(theta + DYNMLP_b1(D, H), H);
+    xavier_init(theta + DYNMLP_W2(D, H), H, D, r);
+    vec_zero(theta + DYNMLP_b2(D, H), D);
+}
+
+void dynmlp_forward(const DynMLP *net, const double *theta,
+                    const double *z, double t, double *out,
+                    Workspace *ws) {
+    int D = net->D, H = net->H;
+    const double *W1 = theta + DYNMLP_W1(D, H);
+    const double *b1 = theta + DYNMLP_b1(D, H);
+    const double *W2 = theta + DYNMLP_W2(D, H);
+    const double *b2 = theta + DYNMLP_b2(D, H);
+
+    double *x     = ws->x;
+    double *h_pre = ws->h_pre;
+    double *h     = ws->h;
+
+    vec_copy(z, x, D);
+    x[D] = t;
+
+    mat_vec(W1, x, h_pre, H, D + 1);
+    vec_add_scaled(h_pre, 1.0, b1, H);
+
+    for (int i = 0; i < H; i++) h[i] = tanh(h_pre[i]);
+
+    mat_vec(W2, h, out, D, H);
+    vec_add_scaled(out, 1.0, b2, D);
+}
+
+void dynmlp_vjp(const DynMLP *net, const double *theta,
+                const double *z, double t, const double *v,
+                double *vjp_z, double *vjp_theta,
+                Workspace *ws) {
+    int D = net->D;
+    int H = net->H;
+    const double *W1 = theta + DYNMLP_W1(D, H);
+    const double *b1 = theta + DYNMLP_b1(D, H);
+    const double *W2 = theta + DYNMLP_W2(D, H);
+    double *dW1 = vjp_theta + DYNMLP_W1(D, H);
+    double *db1 = vjp_theta + DYNMLP_b1(D, H);
+    double *dW2 = vjp_theta + DYNMLP_W2(D, H);
+    double *db2 = vjp_theta + DYNMLP_b2(D, H);
+
+    double *x     = ws->x;
+    double *h_pre = ws->h_pre;
+    double *h     = ws->h;
+    vec_copy(z, x, D);
+    x[D] = t;
+    mat_vec(W1, x, h_pre, H, D + 1);
+    vec_add_scaled(h_pre, 1.0, b1, H);
+
+    for (int i = 0; i < H; i++)
+        h[i] = tanh(h_pre[i]);
+
+    double *dh     = ws->dh;
+    double *dh_pre = ws->dh_pre;
+    double *dx     = ws->dx;
+    vec_zero(dh, H);
+    vec_zero(dx, D + 1);
+
+    mat_vec_T(W2, v, dh, D, H);
+    mat_outer_add(dW2, 1.0, v, h, D, H);
+    vec_add_scaled(db2, 1.0, v, D);
+
+    for (int i = 0; i < H; i++)
+        dh_pre[i] = (1.0 - h[i] * h[i]) * dh[i];
+
+    mat_vec_T(W1, dh_pre, dx, H, D + 1);
+    mat_outer_add(dW1, 1.0, dh_pre, x, H, D + 1);
+    vec_add_scaled(db1, 1.0, dh_pre, H);
+
+    vec_copy(dx, vjp_z, D);
+}