Replace DynMLP with DynNet

1 files changed, 246 insertions, 52 deletions

diff --git a/src/tests.c b/src/tests.c
index 7d42a5c..2775788 100644
--- a/src/tests.c
+++ b/src/tests.c
@@ -1,5 +1,5 @@
 #include "tests.h"
-#include "dynmlp.h"
+#include "dynnet.h"
 #include "ode_solver.h"
 #include "adjoint.h"
 #include "adam.h"
@@ -41,34 +41,42 @@ void test_ode_solver(void) {
       free(r.y); }
 }
 
-void test_dynmlp_gradients(RNG *r) {
-    const int D = 3, H = 8;
-    const double EPS = 1e-7, TOL = 1e-5;
+/* Helper: time_concat → linear → tanh → linear */
+static DynNet *make_dynmlp_net(int D, int H) {
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    return net;
+}
 
-    int np = dynmlp_nparams(D, H);
-    double *theta = vec_alloc(np);
-    double *z = vec_alloc(D);
-    double *v = vec_alloc(D);
+/* Shared VJP gradient check against finite differences */
+static void check_net_gradients(DynNet *net, double *theta, RNG *r,
+                                  const char *label) {
+    const double EPS = 1e-7, TOL = 1e-5;
+    int D  = net->D;
+    int np = net->total_params;
+    double *z     = vec_alloc(D);
+    double *v     = vec_alloc(D);
     double *out_p = vec_alloc(D);
     double *out_m = vec_alloc(D);
+    double *ws    = vec_alloc(net->total_workspace);
 
-    DynMLP net;
-    dynmlp_init(&net, D, H, theta, r);
     for (int i = 0; i < D; i++) z[i] = rng_normal(r);
     for (int i = 0; i < D; i++) v[i] = rng_normal(r);
     double t = rng_normal(r);
 
-    Workspace ws = workspace_alloc(D, H, np);
-
     double *vjp_z     = vec_zeros(D);
     double *vjp_theta = vec_zeros(np);
-    dynmlp_vjp(&net, theta, z, t, v, vjp_z, vjp_theta, &ws);
+    dynnet_vjp(net, theta, z, t, v, vjp_z, vjp_theta, ws);
 
     double *num_vjp_z = vec_alloc(D);
     for (int i = 0; i < D; i++) {
         double zi = z[i];
-        z[i] = zi + EPS; dynmlp_forward(&net, theta, z, t, out_p, &ws);
-        z[i] = zi - EPS; dynmlp_forward(&net, theta, z, t, out_m, &ws);
+        z[i] = zi + EPS; dynnet_forward(net, theta, z, t, out_p, ws);
+        z[i] = zi - EPS; dynnet_forward(net, theta, z, t, out_m, ws);
         z[i] = zi;
         num_vjp_z[i] = (vec_dot(v, out_p, D) - vec_dot(v, out_m, D)) / (2.0 * EPS);
     }
@@ -81,8 +89,8 @@ void test_dynmlp_gradients(RNG *r) {
     double *num_vjp_theta = vec_alloc(np);
     for (int k = 0; k < np; k++) {
         double tk = theta[k];
-        theta[k] = tk + EPS; dynmlp_forward(&net, theta, z, t, out_p, &ws);
-        theta[k] = tk - EPS; dynmlp_forward(&net, theta, z, t, out_m, &ws);
+        theta[k] = tk + EPS; dynnet_forward(net, theta, z, t, out_p, ws);
+        theta[k] = tk - EPS; dynnet_forward(net, theta, z, t, out_m, ws);
         theta[k] = tk;
         num_vjp_theta[k] = (vec_dot(v, out_p, D) - vec_dot(v, out_m, D)) / (2.0 * EPS);
     }
@@ -92,36 +100,48 @@ void test_dynmlp_gradients(RNG *r) {
         if (e > max_err_theta) max_err_theta = e;
     }
 
-    printf("dL/dz:     max_err=%.2e  %s\n", max_err_z, max_err_z < TOL ? "PASS" : "FAIL");
-    printf("dL/dtheta: max_err=%.2e  %s\n", max_err_theta, max_err_theta < TOL ? "PASS" : "FAIL");
+    printf("%s dL/dz:     max_err=%.2e  %s\n", label, max_err_z,     max_err_z     < TOL ? "PASS" : "FAIL");
+    printf("%s dL/dtheta: max_err=%.2e  %s\n", label, max_err_theta, max_err_theta < TOL ? "PASS" : "FAIL");
 
-    workspace_free(&ws);
-    free(theta); free(z); free(v); free(out_p); free(out_m);
+    free(ws); free(z); free(v); free(out_p); free(out_m);
     free(vjp_z); free(vjp_theta); free(num_vjp_z); free(num_vjp_theta);
 }
 
+void test_dynnet_gradients(RNG *r) {
+    const int D = 3, H = 8;
+    DynNet *net = make_dynmlp_net(D, H);
+    int np = net->total_params;
+    double *theta = vec_alloc(np);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "dynnet");
+    dynnet_free(net);
+    free(theta);
+}
+
 void test_adjoint_gradients(RNG *r) {
     const int D = 2, H = 8;
     const double EPS = 1e-5, atol = 1e-7, rtol = 1e-7;
     const double t0 = 0.0, t1 = 1.0;
 
-    int np = dynmlp_nparams(D, H);
+    DynNet *net = make_dynmlp_net(D, H);
+    int np = net->total_params;
     double *theta  = vec_alloc(np);
-    double *z0 = vec_alloc(D);
+    double *z0     = vec_alloc(D);
     double *target = vec_alloc(D);
 
-    DynMLP net;
-    dynmlp_init(&net, D, H, theta, r);
-    for (int i = 0; i < D; i++) z0[i] = rng_normal(r);
+    dynnet_init_params(net, theta, r);
+    for (int i = 0; i < D; i++) z0[i]     = rng_normal(r);
     for (int i = 0; i < D; i++) target[i] = rng_normal(r);
 
-    NeuralODEOutput out = neural_ode_forward_backward(&net, theta, z0, t0, t1,
+    NeuralODEOutput out = neural_ode_forward_backward(net, theta, z0, t0, t1,
                                                       target, atol, rtol, 10);
 
-    Workspace ws = workspace_alloc(D, H, np);
+    double *ws    = vec_alloc(net->total_workspace);
+    double *neg_a = vec_alloc(D);
+    AdjointCtx ac = { net, theta, ws, neg_a };
 
 #define FWD_LOSS(z0_, theta_) ({ \
-    AdjointCtx ac_ = { net, (theta_), D, np, &ws }; \
+    AdjointCtx ac_ = { net, (theta_), ws, neg_a }; \
     ODEResult r_ = ode_solve(neural_ode_rhs, (z0_), t0, t1, NULL, D, atol, rtol, &ac_); \
     double l_ = 0.0; \
     for (int _i = 0; _i < D; _i++) { double _d = r_.y[_i] - target[_i]; l_ += 0.5*_d*_d; } \
@@ -168,11 +188,13 @@ void test_adjoint_gradients(RNG *r) {
     printf("adjoint dL/dz0:    max_rel_err=%.2e  %s\n", rel_z0, rel_z0 < 1e-3 ? "PASS" : "FAIL");
 
 #undef FWD_LOSS
+    (void)ac;
 
-    workspace_free(&ws);
+    free(ws); free(neg_a);
     free(out.z1); free(out.dL_dz0); free(out.dL_dtheta);
     free(num_dL_dtheta); free(num_dL_dz0);
     free(theta); free(z0); free(target);
+    dynnet_free(net);
 }
 
 void test_multi_obs_adjoint(RNG *r) {
@@ -182,23 +204,23 @@ void test_multi_obs_adjoint(RNG *r) {
 
     double times[5] = { 0.0, 0.5, 1.0, 1.5, 2.0 };
 
-    int np = dynmlp_nparams(D, H);
-    double *theta = vec_alloc(np);
-    double *z0 = vec_alloc(D);
+    DynNet *net = make_dynmlp_net(D, H);
+    int np = net->total_params;
+    double *theta   = vec_alloc(np);
+    double *z0      = vec_alloc(D);
     double *targets = vec_alloc(ntimes * D);
 
-    DynMLP net;
-    dynmlp_init(&net, D, H, theta, r);
+    dynnet_init_params(net, theta, r);
     for (int i = 0; i < D; i++) z0[i] = rng_normal(r);
     for (int i = 0; i < ntimes * D; i++) targets[i] = rng_normal(r);
 
     MultiObsNeuralODEOutput out = neural_ode_forward_backward_multi(
-        &net, theta, z0, times, targets, ntimes, atol, rtol);
+        net, theta, z0, times, targets, ntimes, atol, rtol);
 
-    Workspace ws = workspace_alloc(D, H, np);
-    AdjointCtx ac = { net, theta, D, np, &ws };
+    double *ws    = vec_alloc(net->total_workspace);
+    double *neg_a = vec_alloc(D);
+    AdjointCtx ac = { net, theta, ws, neg_a };
 
-    /* Numerical dL/dtheta */
     double *num_dL_dtheta = vec_alloc(np);
     for (int k = 0; k < np; k++) {
         double tk = theta[k];
@@ -237,7 +259,6 @@ void test_multi_obs_adjoint(RNG *r) {
     printf("multi-obs adjoint dL/dtheta: max_rel_err=%.2e  nfe_fwd=%d  nfe_bwd=%d  %s\n",
            rel_theta, out.nfe_forward, out.nfe_backward, rel_theta < 1e-3 ? "PASS" : "FAIL");
 
-    /* Numerical dL/dz0 */
     double *num_dL_dz0 = vec_alloc(D);
     for (int i = 0; i < D; i++) {
         double zi = z0[i];
@@ -276,10 +297,11 @@ void test_multi_obs_adjoint(RNG *r) {
     printf("multi-obs adjoint dL/dz0:    max_rel_err=%.2e  %s\n",
            rel_z0, rel_z0 < 1e-3 ? "PASS" : "FAIL");
 
-    workspace_free(&ws);
+    free(ws); free(neg_a);
     free(out.z_traj); free(out.dL_dz0); free(out.dL_dtheta);
     free(num_dL_dtheta); free(num_dL_dz0);
     free(theta); free(z0); free(targets);
+    dynnet_free(net);
 }
 
 void test_training(RNG *r) {
@@ -288,20 +310,20 @@ void test_training(RNG *r) {
     const double t0 = 0.0, t1 = 1.0;
     const double atol = 1e-4, rtol = 1e-4;
 
-    DynMLP net;
-    int nparams = dynmlp_nparams(D, H);
+    DynNet *net = make_dynmlp_net(D, H);
+    int nparams = net->total_params;
     double *theta = vec_alloc(nparams);
-    dynmlp_init(&net, D, H, theta, r);
+    dynnet_init_params(net, theta, r);
     Adam adam = adam_init(nparams, 1e-3, 0.9, 0.999, 1e-8);
 
-    double **z0s = (double **)xmalloc(N * sizeof(double *));
+    double **z0s     = (double **)xmalloc(N * sizeof(double *));
     double **targets = (double **)xmalloc(N * sizeof(double *));
     for (int i = 0; i < N; i++) {
         double angle = 2.0 * M_PI * rng_uniform(r);
-        z0s[i] = vec_alloc(D);
+        z0s[i]     = vec_alloc(D);
         targets[i] = vec_alloc(D);
-        z0s[i][0] = cos(angle);
-        z0s[i][1] = sin(angle);
+        z0s[i][0]     = cos(angle);
+        z0s[i][1]     = sin(angle);
         targets[i][0] = -z0s[i][1];
         targets[i][1] =  z0s[i][0];
     }
@@ -316,14 +338,15 @@ void test_training(RNG *r) {
             batch_z0[b]  = z0s[idx];
             batch_tgt[b] = targets[idx];
         }
-        TrainStepResult res = train_step(&net, theta, batch_z0, batch_tgt,
+        TrainStepResult res = train_step(net, theta, batch_z0, batch_tgt,
                                          t0, t1, BATCH, &adam, atol, rtol, 10);
         if ((iter + 1) % 50 == 0)
             printf("iter %3d  loss=%.4f  nfe_fwd=%d\n", iter + 1, res.loss, res.nfe_fwd);
     }
 
-    Workspace ws = workspace_alloc(D, H, nparams);
-    AdjointCtx ac = { net, theta, D, nparams, &ws };
+    double *ws    = vec_alloc(net->total_workspace);
+    double *neg_a = vec_alloc(D);
+    AdjointCtx ac = { net, theta, ws, neg_a };
     double final_loss = 0.0;
     for (int i = 0; i < N; i++) {
         ODEResult fwd = ode_solve(neural_ode_rhs, z0s[i], t0, t1, NULL, D, atol, rtol, &ac);
@@ -336,10 +359,181 @@ void test_training(RNG *r) {
     final_loss /= (double)N;
     printf("Loss: %.4f\n", final_loss);
 
-    workspace_free(&ws);
+    free(ws); free(neg_a);
     adam_free(&adam);
     free(batch_z0); free(batch_tgt);
     for (int i = 0; i < N; i++) { free(z0s[i]); free(targets[i]); }
     free(z0s); free(targets);
     free(theta);
+    dynnet_free(net);
+}
+
+/* ---- New gradient tests ---- */
+
+void test_dynnet_deep_gradients(RNG *r) {
+    const int D = 3, H = 8;
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    double *theta = vec_alloc(net->total_params);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "deep(3 hidden)");
+    dynnet_free(net);
+    free(theta);
+}
+
+void test_dynnet_residual_gradients(RNG *r) {
+    const int D = 4, H = 8;
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, D);   /* bring back to D after time_concat adds 1 */
+    dynnet_add_residual_begin(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_add_residual_end(net);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    double *theta = vec_alloc(net->total_params);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "residual");
+    dynnet_free(net);
+    free(theta);
+}
+
+void test_dynnet_layernorm_gradients(RNG *r) {
+    const int D = 3, H = 8;
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_layernorm(net);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    double *theta = vec_alloc(net->total_params);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "layernorm");
+    dynnet_free(net);
+    free(theta);
+}
+
+void test_dynnet_time_inject_gradients(RNG *r) {
+    /* input → linear → tanh → time_concat → linear → tanh → linear */
+    const int D = 3, H = 8;
+    DynNet *net = dynnet_create(D);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    double *theta = vec_alloc(net->total_params);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "time_inject");
+    dynnet_free(net);
+    free(theta);
+}
+
+void test_dynnet_swish_gradients(RNG *r) {
+    const int D = 3, H = 8;
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_swish(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    double *theta = vec_alloc(net->total_params);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "swish");
+    dynnet_free(net);
+    free(theta);
+}
+
+void test_dynnet_softplus_gradients(RNG *r) {
+    const int D = 3, H = 8;
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_softplus(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+    double *theta = vec_alloc(net->total_params);
+    dynnet_init_params(net, theta, r);
+    check_net_gradients(net, theta, r, "softplus");
+    dynnet_free(net);
+    free(theta);
+}
+
+void test_adjoint_deep(RNG *r) {
+    const int D = 2, H = 8;
+    const double EPS = 1e-5, atol = 1e-7, rtol = 1e-7;
+    const double t0 = 0.0, t1 = 1.0;
+
+    DynNet *net = dynnet_create(D);
+    dynnet_add_time_concat(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, H);
+    dynnet_add_tanh(net);
+    dynnet_add_linear(net, D);
+    dynnet_finalize(net);
+
+    int np = net->total_params;
+    double *theta  = vec_alloc(np);
+    double *z0     = vec_alloc(D);
+    double *target = vec_alloc(D);
+
+    dynnet_init_params(net, theta, r);
+    for (int i = 0; i < D; i++) z0[i]     = rng_normal(r);
+    for (int i = 0; i < D; i++) target[i] = rng_normal(r);
+
+    NeuralODEOutput out = neural_ode_forward_backward(net, theta, z0, t0, t1,
+                                                      target, atol, rtol, 10);
+
+    double *ws    = vec_alloc(net->total_workspace);
+    double *neg_a = vec_alloc(D);
+    AdjointCtx ac = { net, theta, ws, neg_a };
+
+#define FWD_LOSS_DEEP(z0_, theta_) ({ \
+    AdjointCtx ac_ = { net, (theta_), ws, neg_a }; \
+    ODEResult r_ = ode_solve(neural_ode_rhs, (z0_), t0, t1, NULL, D, atol, rtol, &ac_); \
+    double l_ = 0.0; \
+    for (int _i = 0; _i < D; _i++) { double _d = r_.y[_i] - target[_i]; l_ += 0.5*_d*_d; } \
+    free(r_.y); l_; \
+})
+
+    double *num_dL_dtheta = vec_alloc(np);
+    for (int k = 0; k < np; k++) {
+        double tk = theta[k];
+        theta[k] = tk + EPS; double lp = FWD_LOSS_DEEP(z0, theta);
+        theta[k] = tk - EPS; double lm = FWD_LOSS_DEEP(z0, theta);
+        theta[k] = tk;
+        num_dL_dtheta[k] = (lp - lm) / (2.0 * EPS);
+    }
+    double max_num = 0.0, max_err = 0.0;
+    for (int k = 0; k < np; k++) {
+        if (fabs(num_dL_dtheta[k]) > max_num) max_num = fabs(num_dL_dtheta[k]);
+        double e = fabs(out.dL_dtheta[k] - num_dL_dtheta[k]);
+        if (e > max_err) max_err = e;
+    }
+    double rel = max_err / (max_num + 1e-8);
+    printf("adjoint_deep dL/dtheta: max_rel_err=%.2e  nfe_fwd=%d  nfe_bwd=%d  %s\n",
+           rel, out.nfe_forward, out.nfe_backward, rel < 1e-3 ? "PASS" : "FAIL");
+
+#undef FWD_LOSS_DEEP
+    (void)ac;
+
+    free(ws); free(neg_a);
+    free(out.z1); free(out.dL_dz0); free(out.dL_dtheta);
+    free(num_dL_dtheta); free(theta); free(z0); free(target);
+    dynnet_free(net);
 }