From 063b20be98747ba0b288545e653ef9489bf14a54 Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Wed, 18 Mar 2026 14:04:47 +0300 Subject: Replace DynMLP with DynNet --- src/adjoint.c | 88 ++++----- src/cnf.c | 145 +++++++------- src/cnf_train.c | 16 +- src/dynmlp.c | 94 --------- src/dynnet.c | 587 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.c | 41 ++-- src/spiral.c | 10 +- src/test_cnf.c | 84 +++++--- src/tests.c | 298 +++++++++++++++++++++++----- src/train.c | 8 +- 10 files changed, 1064 insertions(+), 307 deletions(-) delete mode 100644 src/dynmlp.c create mode 100644 src/dynnet.c (limited to 'src') diff --git a/src/adjoint.c b/src/adjoint.c index df62927..82d63be 100644 --- a/src/adjoint.c +++ b/src/adjoint.c @@ -26,34 +26,26 @@ void neural_ode_rhs(const double *state, double t, const double *params, int dim, double *out, void *ctx) { (void)params; (void)dim; - AdjointCtx *ac = (AdjointCtx *)ctx; - dynmlp_forward(&ac->net, ac->theta, state, t, out, ac->ws); + dynnet_forward(ac->net, ac->theta, state, t, out, ac->ws); } static void adjoint_dynamics(const double *aug_state, double t, const double *params, - int aug_dim, double *aug_out, void *ctx) { + int aug_dim, double *aug_out, void *ctx) { (void)params; (void)aug_dim; - AdjointCtx *ac = (AdjointCtx *)ctx; - int D = ac->state_dim; - int nparams = ac->nparams; + int D = ac->net->D; + int nparams = ac->net->total_params; const double *z = aug_state; const double *a = aug_state + D; - dynmlp_forward(&ac->net, ac->theta, z, t, aug_out, ac->ws); - - double *neg_a = ac->ws->neg_a; - double *vjp_z = ac->ws->vjp_z; - double *vjp_theta = ac->ws->vjp_theta; - vec_zero(vjp_theta, nparams); - for (int i = 0; i < D; i++) neg_a[i] = -a[i]; - - dynmlp_vjp(&ac->net, ac->theta, z, t, neg_a, vjp_z, vjp_theta, ac->ws); + dynnet_forward(ac->net, ac->theta, z, t, aug_out, ac->ws); - vec_copy(vjp_z, aug_out + D, D); - vec_copy(vjp_theta, aug_out + 2 * D, nparams); + for (int i = 0; i < D; i++) ac->neg_a[i] = -a[i]; + vec_zero(aug_out + 2 * D, nparams); + dynnet_vjp(ac->net, ac->theta, z, t, ac->neg_a, + aug_out + D, aug_out + 2 * D, ac->ws); } static void forward_result_free(ForwardResult *fr) { @@ -64,12 +56,13 @@ static void forward_result_free(ForwardResult *fr) { free(fr->z1); } -static ForwardResult forward_solve(const DynMLP *net, const double *theta, +static ForwardResult forward_solve(DynNet *net, const double *theta, const double *z0, double t0, double t1, double atol, double rtol, int num_checkpoints) { int D = net->D; - Workspace ws = workspace_alloc(D, net->H, net->nparams); - AdjointCtx ac = { *net, theta, D, net->nparams, &ws }; + double *ws = vec_alloc(net->total_workspace); + double *neg_a = vec_alloc(D); + AdjointCtx ac = { net, theta, ws, neg_a }; ForwardResult fr; fr.num_checkpoints = num_checkpoints; @@ -97,23 +90,25 @@ static ForwardResult forward_solve(const DynMLP *net, const double *theta, fr.z1 = vec_alloc(D); vec_copy(fr.states[num_checkpoints], fr.z1, D); - workspace_free(&ws); + free(ws); + free(neg_a); return fr; } -static AdjointResult adjoint_solve(const DynMLP *net, const double *theta, +static AdjointResult adjoint_solve(DynNet *net, const double *theta, const ForwardResult *fr, const double *dL_dz1, double atol, double rtol) { - int D = net->D; - int nparams = net->nparams; + int D = net->D; + int nparams = net->total_params; int aug_dim = 2 * D + nparams; double *aug = vec_zeros(aug_dim); vec_copy(fr->z1, aug, D); vec_copy(dL_dz1, aug + D, D); - Workspace ws = workspace_alloc(D, net->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 }; int total_nfe = 0; for (int k = fr->num_checkpoints; k >= 1; k--) { @@ -132,15 +127,16 @@ static AdjointResult adjoint_solve(const DynMLP *net, const double *theta, ar.dL_dz0 = vec_alloc(D); ar.dL_dtheta = vec_alloc(nparams); ar.nfe = total_nfe; - vec_copy(aug + D, ar.dL_dz0, D); - vec_copy(aug + 2 * D, ar.dL_dtheta, nparams); + vec_copy(aug + D, ar.dL_dz0, D); + vec_copy(aug + 2 * D, ar.dL_dtheta, nparams); - workspace_free(&ws); + free(ws); + free(neg_a); free(aug); return ar; } -NeuralODEOutput neural_ode_forward_backward(const DynMLP *net, const double *theta, +NeuralODEOutput neural_ode_forward_backward(DynNet *net, const double *theta, const double *z0, double t0, double t1, const double *target, double atol, double rtol, int num_checkpoints) { @@ -167,7 +163,7 @@ NeuralODEOutput neural_ode_forward_backward(const DynMLP *net, const double *the } static MultiObsAdjointResult adjoint_solve_multi( - const DynMLP *net, + DynNet *net, const double *theta, const double *z_traj, const double *times, @@ -177,16 +173,17 @@ static MultiObsAdjointResult adjoint_solve_multi( double rtol) { int D = net->D; - int nparams = net->nparams; + int nparams = net->total_params; int aug_dim = 2 * D + nparams; - Workspace ws = workspace_alloc(D, net->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 *a = vec_alloc(D); + double *a = vec_alloc(D); double *dtheta = vec_zeros(nparams); - double *aug = vec_alloc(aug_dim); - int total_nfe = 0; + double *aug = vec_alloc(aug_dim); + int total_nfe = 0; vec_copy(dL_dz_each + (ntimes - 1) * D, a, D); @@ -202,8 +199,8 @@ static MultiObsAdjointResult adjoint_solve_multi( total_nfe += seg.nfe; free(seg.y); - vec_copy(aug + D, a, D); - vec_copy(aug + 2 * D, dtheta, nparams); + vec_copy(aug + D, a, D); + vec_copy(aug + 2 * D, dtheta, nparams); /* Kick: add per-observation loss gradient at time t_{i-1} */ vec_add_scaled(a, 1.0, dL_dz_each + (i - 1) * D, D); @@ -214,13 +211,14 @@ static MultiObsAdjointResult adjoint_solve_multi( result.dL_dtheta = dtheta; result.nfe = total_nfe; - workspace_free(&ws); + free(ws); + free(neg_a); free(aug); return result; } MultiObsNeuralODEOutput neural_ode_forward_backward_multi( - const DynMLP *net, + DynNet *net, const double *theta, const double *z0, const double *times, @@ -230,12 +228,14 @@ MultiObsNeuralODEOutput neural_ode_forward_backward_multi( double rtol) { int D = net->D; - Workspace ws = workspace_alloc(D, net->H, net->nparams); - AdjointCtx ac = { *net, theta, D, net->nparams, &ws }; + double *ws = vec_alloc(net->total_workspace); + double *neg_a = vec_alloc(D); + AdjointCtx ac = { net, theta, ws, neg_a }; ODEResult fwd = ode_solve_times(neural_ode_rhs, z0, times, ntimes, NULL, D, atol, rtol, &ac); - workspace_free(&ws); + free(ws); + free(neg_a); double *dL_dz_each = vec_alloc(ntimes * D); for (int i = 0; i < ntimes * D; i++) diff --git a/src/cnf.c b/src/cnf.c index 9b121e6..08d8dd8 100644 --- a/src/cnf.c +++ b/src/cnf.c @@ -4,64 +4,59 @@ /* Compute tr(df/dz) via central finite differences. * Uses z_p (D) as scratch; writes intermediate outputs to f_p and f_m (D). */ -static double compute_trace_fd(const DynMLP *net, const double *theta, - const double *z, double t, Workspace *ws, +static double compute_trace_fd(DynNet *net, const double *theta, + const double *z, double t, double *ws, double eps, double *z_p, double *f_p, double *f_m) { int D = net->D; double tr = 0.0; for (int i = 0; i < D; i++) { vec_copy(z, z_p, D); z_p[i] += eps; - dynmlp_forward(net, theta, z_p, t, f_p, ws); + dynnet_forward(net, theta, z_p, t, f_p, ws); z_p[i] = z[i] - eps; - dynmlp_forward(net, theta, z_p, t, f_m, ws); + dynnet_forward(net, theta, z_p, t, f_m, ws); tr += (f_p[i] - f_m[i]) / (2.0 * eps); } return tr; } -/* ---- Forward augmented ODE ---- */ typedef struct { - DynMLP net; + DynNet *net; const double *theta; int D; - Workspace *ws; + double *ws; double eps; - double *z_p, *f_p, *f_m; /* D */ + double *z_p, *f_p, *f_m; } CNFFwdCtx; static void cnf_fwd_rhs(const double *aug, double t, const double *params, int aug_dim, double *out, void *ctx) { (void)params; (void)aug_dim; - CNFFwdCtx *c = (CNFFwdCtx *)ctx; + CNFFwdCtx *c = (CNFFwdCtx *)ctx; int D = c->D; - dynmlp_forward(&c->net, c->theta, aug, t, out, c->ws); - out[D] = -compute_trace_fd(&c->net, c->theta, aug, t, c->ws, + dynnet_forward(c->net, c->theta, aug, t, out, c->ws); + out[D] = -compute_trace_fd(c->net, c->theta, aug, t, c->ws, c->eps, c->z_p, c->f_p, c->f_m); } -/* aug_state = [z(D), a_z(D), g(nparams)] - * a_logp (= dL/d(logdet)) is a constant stored in ctx since d(logdet)/dt - * does not depend on logdet itself. */ typedef struct { - DynMLP net; + DynNet *net; const double *theta; int D; int nparams; - Workspace *ws; + double *ws; double eps; - double a_logp; - /* scratch buffers */ - double *z_oj; /* D: outer perturbation (trace grad z & theta) */ - double *z_p2, *f_p2, *f_m2; /* D: inner scratch for nested compute_trace_fd */ - double *neg_a_z; /* D */ - double *dg_tmp; /* nparams */ - double *da_z_tmp; /* D (vjp_z scratch, not used as output) */ - double *v_ei; /* D: one-hot for trace grad theta */ - double *vjp_z_tmp; /* D */ - double *vjp_th_p; /* nparams */ - double *vjp_th_m; /* nparams */ + double a_logp; + double *z_oj; + double *z_p2, *f_p2, *f_m2; + double *neg_a_z; + double *dg_tmp; + double *da_z_tmp; + double *v_ei; + double *vjp_z_tmp; + double *vjp_th_p; + double *vjp_th_m; } CNFAdjCtx; static void cnf_adj_rhs(const double *aug, double t, const double *params, @@ -73,23 +68,23 @@ static void cnf_adj_rhs(const double *aug, double t, const double *params, const double *z = aug; const double *a_z = aug + D; - dynmlp_forward(&cc->net, cc->theta, z, t, out, cc->ws); + dynnet_forward(cc->net, cc->theta, z, t, out, cc->ws); for (int i = 0; i < D; i++) cc->neg_a_z[i] = -a_z[i]; vec_zero(cc->dg_tmp, nparams); - dynmlp_vjp(&cc->net, cc->theta, z, t, cc->neg_a_z, + dynnet_vjp(cc->net, cc->theta, z, t, cc->neg_a_z, out + D, cc->dg_tmp, cc->ws); vec_copy(cc->dg_tmp, out + 2 * D, nparams); - if (cc->a_logp == 0.0) return; /* no trace gradient terms needed */ + if (cc->a_logp == 0.0) return; for (int j = 0; j < D; j++) { vec_copy(z, cc->z_oj, D); cc->z_oj[j] += cc->eps; - double tr_p = compute_trace_fd(&cc->net, cc->theta, cc->z_oj, t, cc->ws, + double tr_p = compute_trace_fd(cc->net, cc->theta, cc->z_oj, t, cc->ws, cc->eps, cc->z_p2, cc->f_p2, cc->f_m2); cc->z_oj[j] = z[j] - cc->eps; - double tr_m = compute_trace_fd(&cc->net, cc->theta, cc->z_oj, t, cc->ws, + double tr_m = compute_trace_fd(cc->net, cc->theta, cc->z_oj, t, cc->ws, cc->eps, cc->z_p2, cc->f_p2, cc->f_m2); out[D + j] += cc->a_logp * (tr_p - tr_m) / (2.0 * cc->eps); } @@ -101,12 +96,12 @@ static void cnf_adj_rhs(const double *aug, double t, const double *params, vec_copy(z, cc->z_oj, D); cc->z_oj[i] += cc->eps; vec_zero(cc->vjp_th_p, nparams); - dynmlp_vjp(&cc->net, cc->theta, cc->z_oj, t, cc->v_ei, + dynnet_vjp(cc->net, cc->theta, cc->z_oj, t, cc->v_ei, cc->vjp_z_tmp, cc->vjp_th_p, cc->ws); cc->z_oj[i] = z[i] - cc->eps; vec_zero(cc->vjp_th_m, nparams); - dynmlp_vjp(&cc->net, cc->theta, cc->z_oj, t, cc->v_ei, + dynnet_vjp(cc->net, cc->theta, cc->z_oj, t, cc->v_ei, cc->vjp_z_tmp, cc->vjp_th_m, cc->ws); cc->v_ei[i] = 0.0; @@ -120,27 +115,39 @@ static void cnf_adj_rhs(const double *aug, double t, const double *params, /* ---- Public API ---- */ void cnf_init(CNF *cnf, int D, int H, double *theta, RNG *r) { - dynmlp_init(&cnf->net, D, H, theta, r); - cnf->nparams = dynmlp_nparams(D, 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); + dynnet_init_params(net, theta, r); + cnf->net = net; + cnf->nparams = net->total_params; cnf->trace_eps = 1e-5; } -CNFSampleResult cnf_sample(const CNF *cnf, const double *theta, +void cnf_free(CNF *cnf) { + dynnet_free(cnf->net); + cnf->net = NULL; +} + +CNFSampleResult cnf_sample(CNF *cnf, const double *theta, const double *z0, double t0, double t1, double atol, double rtol) { - int D = cnf->net.D; - Workspace ws = workspace_alloc(D, cnf->net.H, cnf->nparams); + int D = cnf->net->D; + double *ws = vec_alloc(cnf->net->total_workspace); double *z_p = vec_alloc(D), *f_p = vec_alloc(D), *f_m = vec_alloc(D); - CNFFwdCtx ctx = { cnf->net, theta, D, &ws, cnf->trace_eps, z_p, f_p, f_m }; + CNFFwdCtx ctx = { cnf->net, theta, D, ws, cnf->trace_eps, z_p, f_p, f_m }; double *aug0 = vec_zeros(D + 1); vec_copy(z0, aug0, D); - + ODEResult res = ode_solve(cnf_fwd_rhs, aug0, t0, t1, NULL, D + 1, atol, rtol, &ctx); free(aug0); free(z_p); free(f_p); free(f_m); - workspace_free(&ws); + free(ws); CNFSampleResult out; out.z1 = vec_alloc(D); @@ -151,13 +158,13 @@ CNFSampleResult cnf_sample(const CNF *cnf, const double *theta, return out; } -CNFLogProbResult cnf_log_prob(const CNF *cnf, const double *theta, +CNFLogProbResult cnf_log_prob(CNF *cnf, const double *theta, const double *z1, double t0, double t1, double atol, double rtol) { - int D = cnf->net.D; - Workspace ws = workspace_alloc(D, cnf->net.H, cnf->nparams); + int D = cnf->net->D; + double *ws = vec_alloc(cnf->net->total_workspace); double *z_p = vec_alloc(D), *f_p = vec_alloc(D), *f_m = vec_alloc(D); - CNFFwdCtx ctx = { cnf->net, theta, D, &ws, cnf->trace_eps, z_p, f_p, f_m }; + CNFFwdCtx ctx = { cnf->net, theta, D, ws, cnf->trace_eps, z_p, f_p, f_m }; double *aug1 = vec_zeros(D + 1); vec_copy(z1, aug1, D); @@ -165,10 +172,8 @@ CNFLogProbResult cnf_log_prob(const CNF *cnf, const double *theta, NULL, D + 1, atol, rtol, &ctx); free(aug1); free(z_p); free(f_p); free(f_m); - workspace_free(&ws); + free(ws); - /* Backward integral gives ∫_{t1}^{t0} -tr dt = -delta_logp_fwd, - * so delta_logp = log p(z1) - log p(z0) = -res.y[D]. */ CNFLogProbResult out; out.z0 = vec_alloc(D); vec_copy(res.y, out.z0, D); @@ -178,36 +183,36 @@ CNFLogProbResult cnf_log_prob(const CNF *cnf, const double *theta, return out; } -CNFBackwardResult cnf_backward(const CNF *cnf, const double *theta, +CNFBackwardResult cnf_backward(CNF *cnf, const double *theta, const double *z1, double dL_dlogp, const double *dL_dz1_in, double t0, double t1, double atol, double rtol) { - int D = cnf->net.D; + int D = cnf->net->D; int nparams = cnf->nparams; int aug_dim = 2 * D + nparams; - Workspace ws = workspace_alloc(D, cnf->net.H, nparams); + double *ws = vec_alloc(cnf->net->total_workspace); CNFAdjCtx ctx; - ctx.net = cnf->net; - ctx.theta = theta; - ctx.D = D; + ctx.net = cnf->net; + ctx.theta = theta; + ctx.D = D; ctx.nparams = nparams; - ctx.ws = &ws; - ctx.eps = cnf->trace_eps; - ctx.a_logp = dL_dlogp; - ctx.z_oj = vec_alloc(D); - ctx.z_p2 = vec_alloc(D); - ctx.f_p2 = vec_alloc(D); - ctx.f_m2 = vec_alloc(D); + ctx.ws = ws; + ctx.eps = cnf->trace_eps; + ctx.a_logp = dL_dlogp; + ctx.z_oj = vec_alloc(D); + ctx.z_p2 = vec_alloc(D); + ctx.f_p2 = vec_alloc(D); + ctx.f_m2 = vec_alloc(D); ctx.neg_a_z = vec_alloc(D); - ctx.dg_tmp = vec_alloc(nparams); + ctx.dg_tmp = vec_alloc(nparams); ctx.da_z_tmp = vec_alloc(D); - ctx.v_ei = vec_zeros(D); + ctx.v_ei = vec_zeros(D); ctx.vjp_z_tmp = vec_alloc(D); - ctx.vjp_th_p = vec_alloc(nparams); - ctx.vjp_th_m = vec_alloc(nparams); + ctx.vjp_th_p = vec_alloc(nparams); + ctx.vjp_th_m = vec_alloc(nparams); double *aug = vec_zeros(aug_dim); vec_copy(z1, aug, D); @@ -220,14 +225,14 @@ CNFBackwardResult cnf_backward(const CNF *cnf, const double *theta, CNFBackwardResult out; out.dL_dz0 = vec_alloc(D); out.dL_dtheta = vec_alloc(nparams); - vec_copy(res.y + D, out.dL_dz0, D); - vec_copy(res.y + 2 * D, out.dL_dtheta, nparams); + vec_copy(res.y + D, out.dL_dz0, D); + vec_copy(res.y + 2 * D, out.dL_dtheta, nparams); out.nfe = res.nfe; free(res.y); - free(ctx.z_oj); free(ctx.z_p2); free(ctx.f_p2); free(ctx.f_m2); + free(ctx.z_oj); free(ctx.z_p2); free(ctx.f_p2); free(ctx.f_m2); free(ctx.neg_a_z); free(ctx.dg_tmp); free(ctx.da_z_tmp); free(ctx.v_ei); free(ctx.vjp_z_tmp); free(ctx.vjp_th_p); free(ctx.vjp_th_m); - workspace_free(&ws); + free(ws); return out; } diff --git a/src/cnf_train.c b/src/cnf_train.c index 21b357f..d88aa3c 100644 --- a/src/cnf_train.c +++ b/src/cnf_train.c @@ -51,7 +51,17 @@ void cnf_train_demo(RNG *r) { const double atol = 1e-4, rtol = 1e-4; const double LR = 1e-3; - int nparams = dynmlp_nparams(D, H); + int nparams = 0; + { + DynNet *tmp = dynnet_create(D); + dynnet_add_time_concat(tmp); + dynnet_add_linear(tmp, H); + dynnet_add_tanh(tmp); + dynnet_add_linear(tmp, D); + dynnet_finalize(tmp); + nparams = tmp->total_params; + dynnet_free(tmp); + } double *theta = vec_alloc(nparams); CNF cnf; cnf_init(&cnf, D, H, theta, r); @@ -80,8 +90,7 @@ void cnf_train_demo(RNG *r) { double log_pm = log_pb + sr.delta_logp; double log_pt = log_p_target(z1); - * = log_pb + delta_logp - log_pt - * (KL(p_model || p_target) estimator) */ + /* KL(p_model || p_target) estimator */ loss += log_pm - log_pt; double dL_dz1[2]; @@ -124,5 +133,6 @@ void cnf_train_demo(RNG *r) { printf("Total parameters: %d\n", nparams); adam_free(&adam); + cnf_free(&cnf); free(theta); } diff --git a/src/dynmlp.c b/src/dynmlp.c deleted file mode 100644 index 70c9b75..0000000 --- a/src/dynmlp.c +++ /dev/null @@ -1,94 +0,0 @@ -#include "dynmlp.h" - -#include - -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); -} diff --git a/src/dynnet.c b/src/dynnet.c new file mode 100644 index 0000000..0f95ef9 --- /dev/null +++ b/src/dynnet.c @@ -0,0 +1,587 @@ +#include "dynnet.h" + +#include +#include +#include +#include + + +typedef enum { + PRE_LINEAR, PRE_TANH, PRE_SOFTPLUS, PRE_SWISH, PRE_LAYERNORM, + PRE_TIME_CONCAT, PRE_RESIDUAL_BEGIN, PRE_RESIDUAL_END +} PreType; + +typedef struct { int type; int param; } PreLayer; + + +typedef struct { int in_dim; int out_dim; } LinearCfg; +typedef struct { int dim; } ActCfg; +typedef struct { int in_dim; double *t_ptr; } TimeConcatCfg; +typedef struct { DynNet *subnet; int dim; double *t_ptr; } ResidualCfg; + + +static const LayerOps linear_ops; +static const LayerOps tanh_ops; +static const LayerOps softplus_ops; +static const LayerOps swish_ops; +static const LayerOps layernorm_ops; +static const LayerOps time_concat_ops; +static const LayerOps residual_ops; + +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; +} + + +static void linear_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + (void)workspace; + const LinearCfg *c = (const LinearCfg *)cfg; + const double *W = theta; + const double *b = theta + c->in_dim * c->out_dim; + mat_vec(W, in, out, c->out_dim, c->in_dim); + vec_add_scaled(out, 1.0, b, c->out_dim); +} + +static void linear_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + (void)workspace; + const LinearCfg *c = (const LinearCfg *)cfg; + const double *W = theta; + double *dW = v_theta; + double *db = v_theta + c->in_dim * c->out_dim; + /* Compute dW and db before zeroing v_in (safe if v_in aliases in) */ + mat_outer_add(dW, 1.0, v_out, in, c->out_dim, c->in_dim); + vec_add_scaled(db, 1.0, v_out, c->out_dim); + vec_zero(v_in, c->in_dim); + mat_vec_T(W, v_out, v_in, c->out_dim, c->in_dim); +} + +static int linear_nparams(const void *cfg) { + const LinearCfg *c = (const LinearCfg *)cfg; + return c->in_dim * c->out_dim + c->out_dim; +} + +static int linear_ws(const void *cfg) { (void)cfg; return 0; } + + +static void tanh_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + (void)theta; (void)workspace; + const ActCfg *c = (const ActCfg *)cfg; + for (int i = 0; i < c->dim; i++) out[i] = tanh(in[i]); +} + +static void tanh_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + (void)theta; (void)v_theta; (void)workspace; + const ActCfg *c = (const ActCfg *)cfg; + for (int i = 0; i < c->dim; i++) { + double h = tanh(in[i]); + v_in[i] = (1.0 - h * h) * v_out[i]; + } +} + +static int act_nparams(const void *cfg) { (void)cfg; return 0; } +static int act_ws(const void *cfg) { (void)cfg; return 0; } + + +static inline double softplus(double x) { + return x > 0.0 ? x + log1p(exp(-x)) : log1p(exp(x)); +} + +static inline double sigmoid(double x) { + return x > 0.0 ? 1.0 / (1.0 + exp(-x)) : exp(x) / (1.0 + exp(x)); +} + +static void softplus_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + (void)theta; (void)workspace; + const ActCfg *c = (const ActCfg *)cfg; + for (int i = 0; i < c->dim; i++) out[i] = softplus(in[i]); +} + +static void softplus_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + (void)theta; (void)v_theta; (void)workspace; + const ActCfg *c = (const ActCfg *)cfg; + for (int i = 0; i < c->dim; i++) + v_in[i] = sigmoid(in[i]) * v_out[i]; +} + + +static void swish_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + (void)theta; (void)workspace; + const ActCfg *c = (const ActCfg *)cfg; + for (int i = 0; i < c->dim; i++) out[i] = in[i] * sigmoid(in[i]); +} + +static void swish_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + (void)theta; (void)v_theta; (void)workspace; + const ActCfg *c = (const ActCfg *)cfg; + for (int i = 0; i < c->dim; i++) { + double s = sigmoid(in[i]); + v_in[i] = (s + in[i] * s * (1.0 - s)) * v_out[i]; + } +} + + +static int layernorm_ws(const void *cfg) { + const ActCfg *c = (const ActCfg *)cfg; + return 2 + 2 * c->dim; +} + +static int layernorm_nparams(const void *cfg) { + const ActCfg *c = (const ActCfg *)cfg; + return 2 * c->dim; +} + +static void layernorm_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + const ActCfg *c = (const ActCfg *)cfg; + int dim = c->dim; + const double eps = 1e-5; + const double *gamma = theta; + const double *beta = theta + dim; + + double mean = 0.0; + for (int i = 0; i < dim; i++) mean += in[i]; + mean /= (double)dim; + + double var = 0.0; + for (int i = 0; i < dim; i++) { double d = in[i] - mean; var += d * d; } + var /= (double)dim; + + workspace[0] = mean; + workspace[1] = var; + double *x_hat = workspace + 2; + double std = sqrt(var + eps); + + for (int i = 0; i < dim; i++) { + x_hat[i] = (in[i] - mean) / std; + out[i] = gamma[i] * x_hat[i] + beta[i]; + } +} + +static void layernorm_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + (void)in; + const ActCfg *c = (const ActCfg *)cfg; + int dim = c->dim; + const double eps = 1e-5; + const double *gamma = theta; + double *dg = v_theta; + double *db = v_theta + dim; + + double mean = workspace[0]; (void)mean; + double var = workspace[1]; + double std = sqrt(var + eps); + double *x_hat = workspace + 2; + double *dx_hat = workspace + 2 + dim; + + for (int i = 0; i < dim; i++) { + dg[i] += v_out[i] * x_hat[i]; + db[i] += v_out[i]; + dx_hat[i] = v_out[i] * gamma[i]; + } + + double dvar = 0.0, sum_dx = 0.0; + for (int i = 0; i < dim; i++) { + dvar += dx_hat[i] * x_hat[i]; + sum_dx += dx_hat[i]; + } + dvar *= -0.5 / (std * std); + double dmean = -sum_dx / std; + + for (int i = 0; i < dim; i++) { + v_in[i] = dx_hat[i] / std + + dvar * 2.0 * x_hat[i] * std / (double)dim + + dmean / (double)dim; + } +} + + +static void time_concat_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + (void)theta; (void)workspace; + const TimeConcatCfg *c = (const TimeConcatCfg *)cfg; + vec_copy(in, out, c->in_dim); + out[c->in_dim] = *c->t_ptr; +} + +static void time_concat_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + (void)theta; (void)in; (void)v_theta; (void)workspace; + const TimeConcatCfg *c = (const TimeConcatCfg *)cfg; + vec_copy(v_out, v_in, c->in_dim); +} + +static int time_concat_nparams(const void *cfg) { (void)cfg; return 0; } +static int time_concat_ws(const void *cfg) { (void)cfg; return 0; } + + +static void residual_forward(const void *cfg, const double *theta, + const double *in, double *out, double *workspace) { + const ResidualCfg *c = (const ResidualCfg *)cfg; + dynnet_forward(c->subnet, theta, in, *c->t_ptr, out, workspace); + for (int i = 0; i < c->dim; i++) out[i] += in[i]; +} + +static void residual_vjp(const void *cfg, const double *theta, + const double *in, const double *v_out, + double *v_in, double *v_theta, double *workspace) { + const ResidualCfg *c = (const ResidualCfg *)cfg; + dynnet_vjp(c->subnet, theta, in, *c->t_ptr, v_out, v_in, v_theta, workspace); + for (int i = 0; i < c->dim; i++) v_in[i] += v_out[i]; +} + +static int residual_nparams(const void *cfg) { + const ResidualCfg *c = (const ResidualCfg *)cfg; + return c->subnet->total_params; +} + +static int residual_ws(const void *cfg) { + const ResidualCfg *c = (const ResidualCfg *)cfg; + return c->subnet->total_workspace; +} + + +static const LayerOps linear_ops = { linear_forward, linear_vjp, linear_nparams, linear_ws }; +static const LayerOps tanh_ops = { tanh_forward, tanh_vjp, act_nparams, act_ws }; +static const LayerOps softplus_ops = { softplus_forward, softplus_vjp, act_nparams, act_ws }; +static const LayerOps swish_ops = { swish_forward, swish_vjp, act_nparams, act_ws }; +static const LayerOps layernorm_ops = { layernorm_forward, layernorm_vjp, layernorm_nparams, layernorm_ws }; +static const LayerOps time_concat_ops = { time_concat_forward, time_concat_vjp, time_concat_nparams, time_concat_ws }; +static const LayerOps residual_ops = { residual_forward, residual_vjp, residual_nparams, residual_ws }; + + +DynNet *dynnet_create(int D) { + DynNet *net = (DynNet *)xcalloc(1, sizeof(DynNet)); + net->D = D; + net->_cur_dim = D; + return net; +} + +static void pre_append(DynNet *net, int type, int param) { + if (net->_pre_n == net->_pre_cap) { + int new_cap = net->_pre_cap ? net->_pre_cap * 2 : 8; + PreLayer *old = (PreLayer *)net->_pre; + PreLayer *new_buf = (PreLayer *)xmalloc((size_t)new_cap * sizeof(PreLayer)); + if (old) { + memcpy(new_buf, old, (size_t)net->_pre_n * sizeof(PreLayer)); + free(old); + } + net->_pre = new_buf; + net->_pre_cap = new_cap; + } + ((PreLayer *)net->_pre)[net->_pre_n++] = (PreLayer){ type, param }; +} + +void dynnet_add_linear(DynNet *net, int out_dim) { + pre_append(net, PRE_LINEAR, out_dim); + net->_cur_dim = out_dim; +} + +void dynnet_add_tanh(DynNet *net) { + pre_append(net, PRE_TANH, 0); +} + +void dynnet_add_softplus(DynNet *net) { + pre_append(net, PRE_SOFTPLUS, 0); +} + +void dynnet_add_swish(DynNet *net) { + pre_append(net, PRE_SWISH, 0); +} + +void dynnet_add_layernorm(DynNet *net) { + pre_append(net, PRE_LAYERNORM, 0); +} + +void dynnet_add_residual_begin(DynNet *net) { + pre_append(net, PRE_RESIDUAL_BEGIN, 0); +} + +void dynnet_add_residual_end(DynNet *net) { + pre_append(net, PRE_RESIDUAL_END, 0); +} + +void dynnet_add_time_concat(DynNet *net) { + pre_append(net, PRE_TIME_CONCAT, 0); + net->_cur_dim++; +} + + +typedef struct { + const LayerOps *ops; + void *cfg; + int in_dim; + int out_dim; +} LInfo; + +static int find_matching_end(PreLayer *pre, int n, int start) { + int depth = 0; + for (int i = start; i < n; i++) { + if (pre[i].type == PRE_RESIDUAL_BEGIN) depth++; + else if (pre[i].type == PRE_RESIDUAL_END) { + if (--depth == 0) return i; + } + } + return -1; +} + +static LInfo *build_linfo(PreLayer *pre, int n, int D_in, + double *t_ptr, int *out_n, int *out_D) { + int cap = n > 0 ? n : 1; + LInfo *linfo = (LInfo *)xmalloc((size_t)cap * sizeof(LInfo)); + int cnt = 0; + int cur = D_in; + + for (int i = 0; i < n; ) { + if (cnt == cap) { + cap *= 2; + LInfo *tmp = (LInfo *)xmalloc((size_t)cap * sizeof(LInfo)); + memcpy(tmp, linfo, (size_t)cnt * sizeof(LInfo)); + free(linfo); + linfo = tmp; + } + PreLayer p = pre[i]; + + if (p.type == PRE_LINEAR) { + LinearCfg *cfg = (LinearCfg *)xmalloc(sizeof(LinearCfg)); + cfg->in_dim = cur; cfg->out_dim = p.param; + linfo[cnt++] = (LInfo){ &linear_ops, cfg, cur, p.param }; + cur = p.param; i++; + + } else if (p.type == PRE_TANH) { + ActCfg *cfg = (ActCfg *)xmalloc(sizeof(ActCfg)); + cfg->dim = cur; + linfo[cnt++] = (LInfo){ &tanh_ops, cfg, cur, cur }; + i++; + + } else if (p.type == PRE_SOFTPLUS) { + ActCfg *cfg = (ActCfg *)xmalloc(sizeof(ActCfg)); + cfg->dim = cur; + linfo[cnt++] = (LInfo){ &softplus_ops, cfg, cur, cur }; + i++; + + } else if (p.type == PRE_SWISH) { + ActCfg *cfg = (ActCfg *)xmalloc(sizeof(ActCfg)); + cfg->dim = cur; + linfo[cnt++] = (LInfo){ &swish_ops, cfg, cur, cur }; + i++; + + } else if (p.type == PRE_LAYERNORM) { + ActCfg *cfg = (ActCfg *)xmalloc(sizeof(ActCfg)); + cfg->dim = cur; + linfo[cnt++] = (LInfo){ &layernorm_ops, cfg, cur, cur }; + i++; + + } else if (p.type == PRE_TIME_CONCAT) { + TimeConcatCfg *cfg = (TimeConcatCfg *)xmalloc(sizeof(TimeConcatCfg)); + cfg->in_dim = cur; cfg->t_ptr = t_ptr; + linfo[cnt++] = (LInfo){ &time_concat_ops, cfg, cur, cur + 1 }; + cur++; i++; + + } else if (p.type == PRE_RESIDUAL_BEGIN) { + int end = find_matching_end(pre, n, i); + if (end < 0) { + fprintf(stderr, "dynnet: unmatched residual_begin\n"); abort(); + } + DynNet *subnet = (DynNet *)xcalloc(1, sizeof(DynNet)); + subnet->D = cur; + + int sub_n, sub_D; + LInfo *sub_linfo = build_linfo(pre + i + 1, end - i - 1, + cur, &subnet->current_t, &sub_n, &sub_D); + if (sub_D != cur) { + fprintf(stderr, "dynnet: residual in_dim=%d out_dim=%d mismatch\n", cur, sub_D); + abort(); + } + + subnet->num_layers = sub_n; + subnet->ops = (const LayerOps **)xmalloc((size_t)sub_n * sizeof(LayerOps *)); + subnet->layer_configs = (void **)xmalloc((size_t)sub_n * sizeof(void *)); + subnet->param_offsets = (int *)xmalloc((size_t)(sub_n + 1) * sizeof(int)); + subnet->ws_offsets = (int *)xmalloc((size_t)(sub_n + 1) * sizeof(int)); + subnet->in_dims = (int *)xmalloc((size_t)sub_n * sizeof(int)); + + int p_off = 0, w_off = 0, max_d = cur; + for (int k = 0; k < sub_n; k++) { + subnet->ops[k] = sub_linfo[k].ops; + subnet->layer_configs[k] = sub_linfo[k].cfg; + subnet->in_dims[k] = sub_linfo[k].in_dim; + subnet->param_offsets[k] = p_off; + subnet->ws_offsets[k] = w_off; + p_off += sub_linfo[k].ops->nparams(sub_linfo[k].cfg); + w_off += sub_linfo[k].in_dim + sub_linfo[k].ops->workspace_size(sub_linfo[k].cfg); + if (sub_linfo[k].in_dim > max_d) max_d = sub_linfo[k].in_dim; + if (sub_linfo[k].out_dim > max_d) max_d = sub_linfo[k].out_dim; + } + subnet->total_params = p_off; + subnet->v_buf_offset = w_off; + subnet->max_dim = max_d; + subnet->total_workspace = w_off + 2 * max_d; + free(sub_linfo); + + ResidualCfg *cfg = (ResidualCfg *)xmalloc(sizeof(ResidualCfg)); + cfg->subnet = subnet; cfg->dim = cur; cfg->t_ptr = t_ptr; + linfo[cnt++] = (LInfo){ &residual_ops, cfg, cur, cur }; + i = end + 1; + + } else if (p.type == PRE_RESIDUAL_END) { + i++; /* handled by parent */ + + } else { + fprintf(stderr, "dynnet: unknown pre-layer type %d\n", p.type); abort(); + } + } + + *out_n = cnt; + *out_D = cur; + return linfo; +} + +void dynnet_finalize(DynNet *net) { + int out_n, out_D; + LInfo *linfo = build_linfo((PreLayer *)net->_pre, net->_pre_n, + net->D, &net->current_t, &out_n, &out_D); + if (out_D != net->D) { + fprintf(stderr, "dynnet: final dim %d != D %d\n", out_D, net->D); abort(); + } + + net->num_layers = out_n; + net->ops = (const LayerOps **)xmalloc((size_t)out_n * sizeof(LayerOps *)); + net->layer_configs = (void **)xmalloc((size_t)out_n * sizeof(void *)); + net->param_offsets = (int *)xmalloc((size_t)(out_n + 1) * sizeof(int)); + net->ws_offsets = (int *)xmalloc((size_t)(out_n + 1) * sizeof(int)); + net->in_dims = (int *)xmalloc((size_t)out_n * sizeof(int)); + + int p_off = 0, w_off = 0, max_d = net->D; + for (int i = 0; i < out_n; i++) { + net->ops[i] = linfo[i].ops; + net->layer_configs[i] = linfo[i].cfg; + net->in_dims[i] = linfo[i].in_dim; + net->param_offsets[i] = p_off; + net->ws_offsets[i] = w_off; + p_off += linfo[i].ops->nparams(linfo[i].cfg); + w_off += linfo[i].in_dim + linfo[i].ops->workspace_size(linfo[i].cfg); + if (linfo[i].in_dim > max_d) max_d = linfo[i].in_dim; + if (linfo[i].out_dim > max_d) max_d = linfo[i].out_dim; + } + net->total_params = p_off; + net->v_buf_offset = w_off; + net->max_dim = max_d; + net->total_workspace = w_off + 2 * max_d; + + free(linfo); + free(net->_pre); + net->_pre = NULL; + net->_pre_n = net->_pre_cap = 0; +} + + +void dynnet_init_params(DynNet *net, double *theta, RNG *r) { + for (int i = 0; i < net->num_layers; i++) { + double *th = theta + net->param_offsets[i]; + if (net->ops[i] == &linear_ops) { + LinearCfg *c = (LinearCfg *)net->layer_configs[i]; + xavier_init(th, c->in_dim, c->out_dim, r); + vec_zero(th + c->in_dim * c->out_dim, c->out_dim); + } else if (net->ops[i] == &layernorm_ops) { + ActCfg *c = (ActCfg *)net->layer_configs[i]; + for (int j = 0; j < c->dim; j++) th[j] = 1.0; + vec_zero(th + c->dim, c->dim); + } else if (net->ops[i] == &residual_ops) { + ResidualCfg *c = (ResidualCfg *)net->layer_configs[i]; + dynnet_init_params(c->subnet, th, r); + } + } +} + + +void dynnet_free(DynNet *net) { + for (int i = 0; i < net->num_layers; i++) { + if (net->ops[i] == &residual_ops) { + ResidualCfg *c = (ResidualCfg *)net->layer_configs[i]; + dynnet_free(c->subnet); + } + free(net->layer_configs[i]); + } + free((void *)net->ops); + free(net->layer_configs); + free(net->param_offsets); + free(net->ws_offsets); + free(net->in_dims); + if (net->_pre) free(net->_pre); + free(net); +} + + +void dynnet_forward(DynNet *net, const double *theta, + const double *z, double t, double *out, double *workspace) { + net->current_t = t; + int N = net->num_layers; + + vec_copy(z, workspace + net->ws_offsets[0], net->in_dims[0]); + + for (int i = 0; i < N; i++) { + double *in_i = workspace + net->ws_offsets[i]; + double *out_i = (i < N - 1) ? workspace + net->ws_offsets[i + 1] : out; + double *scratch = in_i + net->in_dims[i]; + net->ops[i]->forward(net->layer_configs[i], + theta + net->param_offsets[i], + in_i, out_i, scratch); + } +} + + +void dynnet_vjp(DynNet *net, const double *theta, + const double *z, double t, const double *v, + double *vjp_z, double *vjp_theta, double *workspace) { + net->current_t = t; + int N = net->num_layers; + + vec_copy(z, workspace + net->ws_offsets[0], net->in_dims[0]); + double *throwaway = workspace + net->v_buf_offset; + for (int i = 0; i < N; i++) { + double *in_i = workspace + net->ws_offsets[i]; + double *out_i = (i < N - 1) ? workspace + net->ws_offsets[i + 1] : throwaway; + double *scratch = in_i + net->in_dims[i]; + net->ops[i]->forward(net->layer_configs[i], + theta + net->param_offsets[i], + in_i, out_i, scratch); + } + + double *v_bufs[2] = { + workspace + net->v_buf_offset, + workspace + net->v_buf_offset + net->max_dim + }; + int cur = 0; + vec_copy(v, v_bufs[cur], net->D); + + for (int i = N - 1; i >= 0; i--) { + double *in_i = workspace + net->ws_offsets[i]; + double *scratch = in_i + net->in_dims[i]; + double *v_in_i = (i == 0) ? vjp_z : v_bufs[1 - cur]; + + vec_zero(v_in_i, net->in_dims[i]); + net->ops[i]->vjp(net->layer_configs[i], + theta + net->param_offsets[i], + in_i, v_bufs[cur], v_in_i, + vjp_theta + net->param_offsets[i], + scratch); + if (i > 0) cur ^= 1; + } +} diff --git a/src/main.c b/src/main.c index 291a7bb..470551a 100644 --- a/src/main.c +++ b/src/main.c @@ -1,5 +1,4 @@ -#include "utils.h" -#include "dynmlp.h" +#include "dynnet.h" #include "ode_solver.h" #include "adjoint.h" #include "adam.h" @@ -18,11 +17,20 @@ int main(void) { /* --- sanity checks --- */ test_ode_solver(); - test_dynmlp_gradients(&r); + test_dynnet_gradients(&r); test_adjoint_gradients(&r); test_multi_obs_adjoint(&r); test_training(&r); + printf("\n--- DynNet tests ---\n"); + test_dynnet_deep_gradients(&r); + test_dynnet_residual_gradients(&r); + test_dynnet_layernorm_gradients(&r); + test_dynnet_time_inject_gradients(&r); + test_dynnet_swish_gradients(&r); + test_dynnet_softplus_gradients(&r); + test_adjoint_deep(&r); + printf("\n--- CNF tests ---\n"); test_cnf_trace(&r); test_cnf_invertibility(&r); @@ -46,10 +54,16 @@ int main(void) { Dataset test_ds = generate_spiral_dataset(TEST_N, t0, t1, noise_std, &r); const int D = 2, H = 32; - DynMLP net; - int nparams = dynmlp_nparams(D, 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); + + 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); @@ -67,12 +81,12 @@ int main(void) { batch_tgt[b] = train_ds.target[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_train, rtol_train, 10); if (iter % LOG_EVERY == 0) { - double test_loss = evaluate(&net, theta, &test_ds, + double test_loss = evaluate(net, theta, &test_ds, t0, t1, atol_eval, rtol_eval); printf("%-6d %-12.6f %-12.6f %-10d %-10d\n", iter, res.loss, test_loss, @@ -85,14 +99,15 @@ int main(void) { free(batch_tgt); printf("\n--------------------------------------------------\n"); - double final_test_loss = evaluate(&net, theta, &test_ds, + double final_test_loss = evaluate(net, theta, &test_ds, t0, t1, atol_eval, rtol_eval); printf("Final test loss : %.6f\n", final_test_loss); printf("Total parameters: %d\n", nparams); printf("\nSample predictions:\n"); - 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 }; for (int s = 0; s < 5; s++) { int idx = (int)(rng_next(&r) % (uint64_t)TEST_N); ODEResult fwd = ode_solve(neural_ode_rhs, @@ -104,10 +119,12 @@ int main(void) { test_ds.target[idx][0], test_ds.target[idx][1]); free(fwd.y); } - workspace_free(&ws); + free(ws); + free(neg_a); free(theta); adam_free(&adam); + dynnet_free(net); dataset_free(&train_ds); dataset_free(&test_ds); diff --git a/src/spiral.c b/src/spiral.c index 5d5cc00..e925ea9 100644 --- a/src/spiral.c +++ b/src/spiral.c @@ -52,12 +52,13 @@ void dataset_free(Dataset *ds) { free(ds->target); } -double evaluate(const DynMLP *net, const double *theta, +double evaluate(DynNet *net, const double *theta, const Dataset *ds, double t0, double t1, double atol, double rtol) { int D = net->D; - Workspace ws = workspace_alloc(D, net->H, net->nparams); - AdjointCtx ac = { *net, theta, D, net->nparams, &ws }; + double *ws = vec_alloc(net->total_workspace); + double *neg_a = vec_alloc(D); + AdjointCtx ac = { net, theta, ws, neg_a }; double total_loss = 0.0; for (int i = 0; i < ds->num_samples; i++) { ODEResult fwd = ode_solve(neural_ode_rhs, ds->z0[i], t0, t1, @@ -68,6 +69,7 @@ double evaluate(const DynMLP *net, const double *theta, } free(fwd.y); } - workspace_free(&ws); + free(ws); + free(neg_a); return total_loss / (double)ds->num_samples; } diff --git a/src/test_cnf.c b/src/test_cnf.c index da84869..f01cb27 100644 --- a/src/test_cnf.c +++ b/src/test_cnf.c @@ -7,45 +7,54 @@ #include #include -/* ---- Test 1: Trace via FD agrees with exact trace for linear f ---- */ +/* ---- Test 1: Trace via FD agrees with exact trace ---- */ void test_cnf_trace(RNG *r) { const int D = 3, H = 8; const double atol = 1e-8, rtol = 1e-8; const double EPS_FD = 1e-5, TOL = 1e-4; - int nparams = dynmlp_nparams(D, H); + CNF cnf; + int nparams = 0; + { + DynNet *tmp = dynnet_create(D); + dynnet_add_time_concat(tmp); + dynnet_add_linear(tmp, H); + dynnet_add_tanh(tmp); + dynnet_add_linear(tmp, D); + dynnet_finalize(tmp); + nparams = tmp->total_params; + dynnet_free(tmp); + } double *theta = vec_alloc(nparams); - DynMLP net; - dynmlp_init(&net, D, H, theta, r); + cnf_init(&cnf, D, H, theta, r); double z[3]; for (int i = 0; i < D; i++) z[i] = rng_normal(r); double t = 0.5; - Workspace ws = workspace_alloc(D, H, nparams); + double *ws = vec_alloc(cnf.net->total_workspace); double *z_p = vec_alloc(D), *f_p = vec_alloc(D), *f_m = vec_alloc(D); double tr_fd = 0.0; for (int i = 0; i < D; i++) { vec_copy(z, z_p, D); z_p[i] += EPS_FD; - dynmlp_forward(&net, theta, z_p, t, f_p, &ws); + dynnet_forward(cnf.net, theta, z_p, t, f_p, ws); z_p[i] = z[i] - EPS_FD; - dynmlp_forward(&net, theta, z_p, t, f_m, &ws); + dynnet_forward(cnf.net, theta, z_p, t, f_m, ws); tr_fd += (f_p[i] - f_m[i]) / (2.0 * EPS_FD); } double *out0 = vec_alloc(D); - dynmlp_forward(&net, theta, z, t, out0, &ws); + dynnet_forward(cnf.net, theta, z, t, out0, ws); double tr_jac = 0.0; for (int i = 0; i < D; i++) { double *col_p = vec_alloc(D), *col_m = vec_alloc(D); vec_copy(z, z_p, D); z_p[i] += EPS_FD; - dynmlp_forward(&net, theta, z_p, t, col_p, &ws); + dynnet_forward(cnf.net, theta, z_p, t, col_p, ws); z_p[i] = z[i] - EPS_FD; - dynmlp_forward(&net, theta, z_p, t, col_m, &ws); - /* J[:,i] = (col_p - col_m) / (2*eps); diagonal entry = row i */ + dynnet_forward(cnf.net, theta, z_p, t, col_m, ws); tr_jac += (col_p[i] - col_m[i]) / (2.0 * EPS_FD); free(col_p); free(col_m); } @@ -54,11 +63,6 @@ void test_cnf_trace(RNG *r) { printf("CNF trace vs full Jacobian trace: err=%.2e tr_fd=%.6f tr_jac=%.6f %s\n", err, tr_fd, tr_jac, err < TOL ? "PASS" : "FAIL"); - CNF cnf; - cnf.net = net; - cnf.nparams = nparams; - cnf.trace_eps = EPS_FD; - double dt = 1e-4; CNFSampleResult sr = cnf_sample(&cnf, theta, z, 0.0, dt, atol, rtol); double delta_approx = -tr_fd * dt; @@ -67,8 +71,9 @@ void test_cnf_trace(RNG *r) { err2, err2 < 1e-2 ? "PASS" : "FAIL"); free(sr.z1); - workspace_free(&ws); - free(theta); free(z_p); free(f_p); free(f_m); free(out0); + free(ws); free(z_p); free(f_p); free(f_m); free(out0); + cnf_free(&cnf); + free(theta); } /* ---- Test 2: cnf_sample followed by cnf_log_prob recovers z0 ---- */ @@ -77,7 +82,17 @@ void test_cnf_invertibility(RNG *r) { const double atol = 1e-8, rtol = 1e-8, TOL = 1e-4; const double t0 = 0.0, t1 = 1.0; - int nparams = dynmlp_nparams(D, H); + int nparams = 0; + { + DynNet *tmp = dynnet_create(D); + dynnet_add_time_concat(tmp); + dynnet_add_linear(tmp, H); + dynnet_add_tanh(tmp); + dynnet_add_linear(tmp, D); + dynnet_finalize(tmp); + nparams = tmp->total_params; + dynnet_free(tmp); + } double *theta = vec_alloc(nparams); CNF cnf; cnf_init(&cnf, D, H, theta, r); @@ -86,7 +101,6 @@ void test_cnf_invertibility(RNG *r) { for (int i = 0; i < D; i++) z0[i] = rng_normal(r); CNFSampleResult sr = cnf_sample(&cnf, theta, z0, t0, t1, atol, rtol); - CNFLogProbResult lr = cnf_log_prob(&cnf, theta, sr.z1, t0, t1, atol, rtol); double err = 0.0; @@ -95,7 +109,6 @@ void test_cnf_invertibility(RNG *r) { err += d * d; } err = sqrt(err); - double logp_err = fabs(sr.delta_logp - lr.delta_logp); printf("CNF invertibility: z0_err=%.2e logp_err=%.2e nfe_fwd=%d nfe_bwd=%d %s\n", @@ -103,6 +116,7 @@ void test_cnf_invertibility(RNG *r) { (err < TOL && logp_err < TOL) ? "PASS" : "FAIL"); free(theta); free(sr.z1); free(lr.z0); + cnf_free(&cnf); } /* ---- Test 3: adjoint gradients via finite differences ---- */ @@ -111,7 +125,17 @@ void test_cnf_gradients(RNG *r) { const double EPS = 1e-5, atol = 1e-7, rtol = 1e-7; const double t0 = 0.0, t1 = 0.5; - int nparams = dynmlp_nparams(D, H); + int nparams = 0; + { + DynNet *tmp = dynnet_create(D); + dynnet_add_time_concat(tmp); + dynnet_add_linear(tmp, H); + dynnet_add_tanh(tmp); + dynnet_add_linear(tmp, D); + dynnet_finalize(tmp); + nparams = tmp->total_params; + dynnet_free(tmp); + } double *theta = vec_alloc(nparams); CNF cnf; cnf_init(&cnf, D, H, theta, r); @@ -164,6 +188,7 @@ void test_cnf_gradients(RNG *r) { free(theta); free(z1); free(dL_dz1); free(br.dL_dz0); free(br.dL_dtheta); free(num_grad); + cnf_free(&cnf); } /* ---- Test 4: training loss decreases on a simple target ---- */ @@ -173,7 +198,17 @@ void test_cnf_training(RNG *r) { const double t0 = 0.0, t1 = 1.0; const double atol = 1e-4, rtol = 1e-4; - int nparams = dynmlp_nparams(D, H); + int nparams = 0; + { + DynNet *tmp = dynnet_create(D); + dynnet_add_time_concat(tmp); + dynnet_add_linear(tmp, H); + dynnet_add_tanh(tmp); + dynnet_add_linear(tmp, D); + dynnet_finalize(tmp); + nparams = tmp->total_params; + dynnet_free(tmp); + } double *theta = vec_alloc(nparams); CNF cnf; cnf_init(&cnf, D, H, theta, r); @@ -225,7 +260,7 @@ void test_cnf_training(RNG *r) { free(z1); } - if (iter == 0) first_loss = loss / BATCH; + if (iter == 0) first_loss = loss / BATCH; if (iter == ITERS - 1) last_loss = loss / BATCH; adam_update(&adam, theta, dL_dtheta_acc); @@ -236,5 +271,6 @@ void test_cnf_training(RNG *r) { first_loss, last_loss, last_loss < first_loss ? "PASS" : "FAIL"); adam_free(&adam); + cnf_free(&cnf); free(theta); } 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); } diff --git a/src/train.c b/src/train.c index ebaf7cd..be6ec8f 100644 --- a/src/train.c +++ b/src/train.c @@ -4,7 +4,7 @@ #include -static double train_one(const DynMLP *net, const double *theta, +static double train_one(DynNet *net, const double *theta, const double *z0, double t0, double t1, const double *target, double *grad_accum, double atol, double rtol, int num_checkpoints, @@ -17,7 +17,7 @@ static double train_one(const DynMLP *net, const double *theta, double d = out.z1[i] - target[i]; loss += 0.5 * d * d; } - for (int i = 0; i < net->nparams; i++) grad_accum[i] += out.dL_dtheta[i]; + for (int i = 0; i < net->total_params; i++) grad_accum[i] += out.dL_dtheta[i]; *nfe_fwd += out.nfe_forward; *nfe_bwd += out.nfe_backward; free(out.z1); @@ -26,11 +26,11 @@ static double train_one(const DynMLP *net, const double *theta, return loss; } -TrainStepResult train_step(const DynMLP *net, double *theta, +TrainStepResult train_step(DynNet *net, double *theta, const double **z0s, const double **targets, double t0, double t1, int batch_size, Adam *adam, double atol, double rtol, int num_checkpoints) { - int nparams = net->nparams; + int nparams = net->total_params; double *grad_accum = vec_zeros(nparams); TrainStepResult res = { 0.0, 0, 0 }; -- cgit v1.2.3