aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/adjoint.c88
-rw-r--r--src/cnf.c145
-rw-r--r--src/cnf_train.c16
-rw-r--r--src/dynmlp.c94
-rw-r--r--src/dynnet.c587
-rw-r--r--src/main.c41
-rw-r--r--src/spiral.c10
-rw-r--r--src/test_cnf.c84
-rw-r--r--src/tests.c298
-rw-r--r--src/train.c8
10 files changed, 1064 insertions, 307 deletions
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 <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);
-}
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 <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+
+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 <stdlib.h>
#include <math.h>
-/* ---- 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 <stdlib.h>
-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 };
generated by cgit v1.2.3 (git 2.51.0) at 2026-09-01 21:25:28 +0000