BUG: .data_set() function does not work with vectors

[ntolley]

Fairly easy to reproduce, if you try to pass a vector for data_set() you get the following:

import jaxley as jx
import numpy as np
from jaxley.channels import Na

cell = jx.Cell()
cell.insert(Na())
net = jx.Network([cell for _ in range(5)])
param_state = None
param_state = net.data_set('Na_gNa', np.repeat(1.0, 5), param_state)
net.record('v')
s = jx.integrate(net, t_max=10, param_state=param_state)

ValueError: Incompatible shapes for broadcasting: (5, 1) and requested shape (1, 5)

Seems like a simple enough change, will open a PR to try and fix this

[ntolley]

Also I should note this isn't fixed by reshaping the input vector

[michaeldeistler]

Thanks! @kyralianaka aren't you using something like this?

[kyralianaka]

Yes I wrote about this a while ago (#430) , but then I just switched to using select and make_trainable when parameter sharing was made possible that way. I agree that this would be nice to fix if we want to keep supporting data_set().