Write trainables to module

jaxleyverse / jaxley

Differentiable neuron simulations with biophysical detail on CPU, GPU, or TPU.

https://jaxley.readthedocs.io

Apache License 2.0

46 stars 7 forks source link

Write trainables to module #470

Closed michaeldeistler closed 1 week ago

michaeldeistler commented 1 week ago

This PR introduces a write_trainables() which adds the trainables to the .nodes and .edges such that they can be used for plotting, printing,...

net.make_trainable("radius")
params = net.get_parameters()

net.write_trainables(params)

Closes #356