Query Synaptic Indicies Inside of .to_jax() instead of .step()

[alexpejovic]

Following are some PR notes. Note the first note, this PR isn't ready to be merged until benchmarks have been run. The purpose of the early PR is to get code feedback.

• This should PR is not ready until I get access to the cluster again.
  • I will comment on the PR again when benchmarks have been run.
• I can't think of a better name for is_integratable, but I recognize it is bad,
  I am open to suggestions!
• All of to_jax() currently cannot be wrapped by if not self.is_integratable,
  since the line values = values.at[self.base.nodes.index.to_numpy()].set(value)
  inside of to_jax() ends up raising an UnexpectedTracerError when jitted.
  • If we want all of to_jax() to be wrapped inside an if condition, we have to
    find a way to avoid the TracerError.
• Line 3475 in modules/base.py is to ensure that test test_view_attrs passes.

[michaeldeistler]

self. is_integratable -> self.to_jax_has_been_run

[alexpejovic]

self. is_integratable -> self.to_jax_has_been_run

This would be great if all of .to_jax() was skipped if self.is_integratable were True, but unfortunately in the current code it is not. See point 3 of my inital comment...

* All of `to_jax()` currently cannot be wrapped by `if not self.is_integratable`,
  since the line `values = values.at[self.base.nodes.index.to_numpy()].set(value)`
  inside of `to_jax()` ends up raising an `UnexpectedTracerError` when jitted.
  
  * If we want all of to_jax() to be wrapped inside an if condition, we have to
    find a way to avoid the `TracerError`.

Maybe self.has_gathered_synapses unless the problem becomes fixed?