`TracerArrayConversionError` in `data_set`

[jnsbck]

The following seems to have broken between 0.8.2 and 0.9.0, presumably due to #606.

from jaxley.synapses import IonotropicSynapse
from jaxley.connect import connect
import jaxley as jx
from jax import jit
import jax.numpy as jnp

net = jx.Network([jx.Cell()]*2)
net.record()
connect(net.cell(0), net.cell(1), IonotropicSynapse())

@jit 
def simulate():
    pstate = net.select(edges=0).data_set("IonotropicSynapse_gS", jnp.array([1.0]), None)
    v = jx.integrate(net, param_state=pstate, t_max=1.0)
    return v

v = simulate()

Now raises TracerArrayConversionError, I am guessing due to

jaxley/jaxley/modules/base.py

Line 1611 in c6d11cd

synapse_inds = (synapse_inds.astype(int) - 1).to_numpy()

EDIT:
I ran a git bisect and #606 seems to have indeed introduced the bug.

commit 08d68b04f9b64f143f517bf3960d9da1d1d98526 (refs/bisect/bad)
Author: Michael Deistler <michael.deistler@uni-tuebingen.de>
Date:   Mon May 19 13:32:22 2025 +0200

    Fix data set indices (#606)
    
    * Make indices be (N, 1) instead of (1, N)
    
    * Make default value of param_state = None
    
    * Add test based off ntolley:data_set_vector
    
    ---------
    
    Co-authored-by: Chase King <iamchaseking@gmail.com>

diff --git a/jaxley/modules/base.py b/jaxley/modules/base.py
index 7f41da4..66a84a0 100644
--- a/jaxley/modules/base.py
+++ b/jaxley/modules/base.py
@@ -1095,7 +1095,7 @@ class Module(ABC):
         self,
         key: str,
         val: Union[float, jnp.ndarray],
-        param_state: Optional[List[Dict]],
+        param_state: Optional[List[Dict]] = None,
     ):
         """Set parameter of module (or its view) to a new value within `jit`.
 
@@ -1112,11 +1112,13 @@ class Module(ABC):
         viewed_inds = self._nodes_in_view if is_node_param else self._edges_in_view
         if key in data.columns:
             not_nan = ~data[key].isna()
+            indices = jnp.asarray(viewed_inds[not_nan]).reshape(-1, 1)  # shape (n_comp, 1)
+            val = jnp.broadcast_to(jnp.asarray(val), (indices.shape[0],))  # shape (n_comp,)
             added_param_state = [
                 {
-                    "indices": np.atleast_2d(viewed_inds[not_nan]),
+                    "indices": indices,
                     "key": key,
-                    "val": jnp.atleast_1d(jnp.asarray(val)),
+                    "val": val,
                 }
             ]
             if param_state is not None:

[Fololo]

I am encountering the same error ...

[jnsbck]

I spend some time looking into this.

The following also does not work:

comp = jx.Compartment()
comp.record()

@jit
def simulate():
    pstate = comp.select(jnp.array([0])).data_set("radius", jnp.array([1.0]), None)
    v = jx.integrate(comp, param_state=pstate, t_max=1.0)
    return v

simulate()

The core issues seems to be seems to be that dataframes cannot be indexed by jax tracers (jax arrays are fine though) since at some point they are getting converted to numpy arrays. I.e. the following is not working

df = pd.DataFrame({"key": [1]})
test = lambda: df.loc[jnp.array([0])]

test() # works
jit(test)() # breaks

This is related to #654. While we could in principle allow indexing with arrays to fix this, fundamentally it will not work for jax tracers. I think the way to fix this is by pushing #646

[Fololo]

Hey @jnsbck do you have any hot-fixes for this?

[jnsbck]

TLDR: Unfortunately not.

Merging #658 would allow jax arrays as indices, but this will not fix the TracerConversionError, since as stated, the fundamental problem is caused by pandas dataframes not being indexable by jax tracers (which are used to determine the current View). Hence the new backend (see #646). @michaeldeistler and I have settled now on how we are going to set this up and I have started implementing a basic version of this to figure out the details. Feel free to join the discussion / follow the PR :)

Let me know if merging #658 would be helpful in the meantime.