Sequencing side-effects in JAX is a known issue in JAX. I have tried to make a custom operational env/task where the step update needs to be sequential.
JAX docs say one needs to tokenize the functions to force the sequence. However, I could not find any task in EvoJax with the tokens. Did anybody ever try that? Or I should think more and redesign my problem to make it more JAX compatible since I use a bunch of jax.lax.cond (constraints are quite complicated for the problem I have).
Sequencing side-effects in JAX is a known issue in JAX. I have tried to make a custom operational env/task where the step update needs to be sequential.
JAX docs say one needs to tokenize the functions to force the sequence. However, I could not find any task in EvoJax with the tokens. Did anybody ever try that? Or I should think more and redesign my problem to make it more JAX compatible since I use a bunch of
jax.lax.cond(constraints are quite complicated for the problem I have).