`get_parameters` doesn't return consistent values after updating parameters

[Kartik-Sama]

I am updating initial values of trainables, if I check the values in the dataframe view values change but get_parameters method still returns old values.
(PS: Asking because I use get_parameters to fetch parameters that are sent into optax transform, but my initial values don't seem to be updated as expected.)

>>> branch.nodes['radius'] 
0    1.0
1    1.0
Name: radius, dtype: float64
>>> branch.get_parameters()
[{'radius': Array([1., 1.], dtype=float32)}]
>>> branch.write_trainables([{"radius":jnp.array([0.25, 1.23], jnp.float32)}])
>>> branch.nodes['radius']
0    0.25
1    1.23
Name: radius, dtype: float64
>>> branch.get_parameters()
[{'radius': Array([1., 1.], dtype=float32)}]

[michaeldeistler]

Hi Kartik,

thanks for reaching out! Unfortunately, I cannot reproduce the issue. The following works for me:

import jaxley as jx
import jax.numpy as jnp

cell = jx.Cell()
cell.make_trainable("radius")
params = cell.get_parameters()
cell.write_trainables([{"radius": jnp.array([0.25], jnp.float32)}])
print(cell.nodes["radius"])

Does this work? If yes, could you provide a full reproducible example of a case that breaks? Thanks!
Michael

[Kartik-Sama]

Hi Michael,

Thanks for the follow up. The <Module>.nodes still works, I also see the updated value when I call cell.nodes["radius"] in your example. However when I call cell.get_parameters() after writing the trainable I still see the old values.

Added the line in your code snippet.

import jaxley as jx
import jax.numpy as jnp

cell = jx.Cell()
cell.make_trainable("radius")
params = cell.get_parameters()
cell.write_trainables([{"radius": jnp.array([0.25], jnp.float32)}])
print(cell.nodes["radius"])
print(cell.get_parameters())