diff options
Diffstat (limited to 'include/adjoint.h')
| -rw-r--r-- | include/adjoint.h | 14 |
1 files changed, 6 insertions, 8 deletions
diff --git a/include/adjoint.h b/include/adjoint.h index d9e71a9..f9fa37a 100644 --- a/include/adjoint.h +++ b/include/adjoint.h @@ -1,15 +1,13 @@ #pragma once -#include "utils.h" -#include "dynmlp.h" +#include "dynnet.h" #include "ode_solver.h" typedef struct { - DynMLP net; + DynNet *net; const double *theta; - int state_dim; - int nparams; - Workspace *ws; + double *ws; /* size net->total_workspace */ + double *neg_a; /* size net->D */ } AdjointCtx; typedef struct { @@ -32,13 +30,13 @@ void neural_ode_rhs(const double *state, double t, const double *params, int dim, double *out, void *ctx); NeuralODEOutput neural_ode_forward_backward( - const DynMLP *net, const double *theta, + DynNet *net, const double *theta, const double *z0, double t0, double t1, const double *target, double atol, double rtol, int num_checkpoints); MultiObsNeuralODEOutput neural_ode_forward_backward_multi( - const DynMLP *net, const double *theta, + DynNet *net, const double *theta, const double *z0, const double *times, const double *targets, int ntimes, double atol, double rtol); |