Get all module parameters (from small param list) without side effects

[bantin]

Hi Jaxley team,

I’m working on smoothing/inference for multicompartment models. I'd like to JIT/grad a function that:

1. Accepts the list of trainable params returned by cell.get_parameters(), then reconstructs the dense all_params dict that step_fn expects.

2. Does so without mutating the original cell or having side effects that will mess with JIT/grad

Obviously this will need to happen using a closure over the jaxley module, but even with that in place I couldn't quite figure out the right way to do it. Right now I’ve implemented an apply_trainable_params helper that mirrors the scatter logic in Module.get_all_parameters: it caches the dense dict produced at build time, along with the indices_set_by_trainables, and inside the jitted function it simply unrolls over that metadata and runs params[key].at[inds].set(values[:, None]) for each trainable entry. That works, but I’d rather reuse an official utility if there is one.

Is there an existing pattern or helper in core Jaxley for going from the trainable list + indices back to the full parameter/state dictionaries in a side-effect free way? It seems like this probably exists, I just wasn't sure of the right pattern here. I'm happy to send along my full code if helpful.

Thanks!
Ben

[bantin]

edit: sorry after looking over things further, it seems like I need to use a closure over the cell and the trainable indices, and then I'll be able to use params_to_pstate and module.get_all_parameters from within JIT. That should have been obvious. So feel free to ignore. I'll close this for now :)