fix several bugs (make it work in the first place

jaxley/integrate.py

@@ -46,18 +46,18 @@ def integrate(
     # At least one stimulus was inserted.
     if module.currents is not None or data_stimuli is not None:
         if module.currents is not None:
-            i_current = module.currents.T
+            i_current = module.currents  # Shape `(num_stimuli, time)`.
             i_inds = module.current_inds.comp_index.to_numpy()
 
             if data_stimuli is not None:
                 # Append stimuli from `data_stimuli`.
-                i_current = jnp.concatenate(
-                    [i_current, jnp.expand_dims(data_stimuli[0], axis=0)]
-                )
-                i_inds = np.concatenate([i_inds, data_stimuli[1].comp_index.to_numpy()])
+                i_current = jnp.concatenate([i_current, data_stimuli[0]])
+                i_inds = jnp.concatenate([i_inds, data_stimuli[1].comp_index.to_numpy()])
         else:
-            i_current = data_stimuli[0]
+            i_current = data_stimuli[0]  # Shape `(num_stimuli, time)`
             i_inds = data_stimuli[1].comp_index.to_numpy()
+
+        i_current = i_current.T  # Shape `(time, num_stimuli)`.
     else:
         # No stimulus was inserted.
         i_current = jnp.asarray([[]]).astype("int32")

jaxley/modules/base.py

@@ -793,7 +793,7 @@ def data_stimulate(
     ):
         """Insert a stimulus into the module within jit (or grad)."""
         nodes = self.set_global_index_and_index(self.view)
-        self.pointer._data_stimulate(current, data_stimuli, nodes)
+        return self.pointer._data_stimulate(current, data_stimuli, nodes)
 
     def set(self, key: str, val: float):
         """Set parameters of the pointer."""