Initial states are trainable

jaxley/integrate.py

@@ -104,9 +104,18 @@ def _body_fun(state, i_stim):
     for channel in module.channels:
         for channel_states in list(channel.channel_states.keys()):
             states[channel_states] = module.jaxnodes[channel_states]
+
+    # Override with the initial states set by `.make_trainable()`.
+    for inds, set_param in zip(module.indices_set_by_trainables, params):
+        for key in set_param.keys():
+            if key in list(states.keys()):  # Only initial states, not parameters.
+                states[key] = states[key].at[inds].set(set_param[key])
+
+    # Write synaptic states. TODO move above when new interface for synapses.
     for key in module.syn_states:
         states[key] = module.syn_states[key]
 
+    # Run simulation.
     _, recordings = nested_checkpoint_scan(
         _body_fun, states, i_current, length=length, nested_lengths=checkpoint_lengths
     )

jaxley/modules/base.py

@@ -319,7 +319,8 @@ def get_all_parameters(self, trainable_params):
         # Override with those parameters set by `.make_trainable()`.
         for inds, set_param in zip(self.indices_set_by_trainables, trainable_params):
             for key in set_param.keys():
-                params[key] = params[key].at[inds].set(set_param[key])
+                if key in list(params.keys()):  # Only parameters, not initial states.
+                    params[key] = params[key].at[inds].set(set_param[key])
 
         # Compute conductance params and append them.
         cond_params = self.init_conds(params)