WIP: Allow vectors to be used with .data_set()

CHANGELOG.md

@@ -30,6 +30,8 @@ net.vis()
 ### Bug fixes
 - Fixed inconsistency with *type* assertions arising due to `numpy` functions returning different `dtypes` on platforms like Windows (#567, @Kartik-Sama)
 
+- Fixed bug where `.data_set()` could not use vectors as input (#574, @ntolley)
+
 # 0.5.0
 
 ### API changes

jaxley/modules/base.py

@@ -1316,9 +1316,9 @@ def get_all_parameters(
             if key in params:  # Only parameters, not initial states.
                 # `inds` is of shape `(num_params, num_comps_per_param)`.
                 # `set_param` is of shape `(num_params,)`
-                # We need to unsqueeze `set_param` to make it `(num_params, 1)` for the
-                # `.set()` to work. This is done with `[:, None]`.
-                params[key] = params[key].at[inds].set(set_param[:, None])
+                # We need to unsqueeze `set_param` to make it `(1, num_params)` for the
+                # `.set()` to work. This is done with `.reshape(1, -1)`.
+                params[key] = params[key].at[inds].set(set_param.reshape(1, -1))
 
         # Compute conductance params and add them to the params dictionary.
         params["axial_conductances"] = self.base._compute_axial_conductances(

tests/test_make_trainable.py

@@ -445,6 +445,16 @@ def test_data_set_vs_make_trainable_network(SimpleNet):
     assert np.max(np.abs(voltages1 - voltages2)) < 1e-8
 
 
+def test_data_set_vector(SimpleNet):
+    net = SimpleNet(2, 2, 4)
+    net.set("radius", np.repeat(1.0, 16))
+
+    param_state = None
+    param_state = net.data_set("length", np.repeat(1.0, 16), param_state)
+    net.record("v")
+    jx.integrate(net, t_max=10, param_state=param_state)
+
+
 def test_make_states_trainable_api(SimpleNet):
     net = SimpleNet(2, 2, 4)
     net.insert(HH())