This repository was archived by the owner on Aug 29, 2025. It is now read-only.
This repository was archived by the owner on Aug 29, 2025. It is now read-only.
[Discussion] Sequencing side-effects in JAX #69
Description
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).