patrick-kidger / patrick-kidger/diffrax
os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false' no longer working with Jax 0.6.0
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
Hello
I had experienced some slowdown in runtime using diffrax with Jax 0.5.3 and was recommended using this flag
os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false' no longer working with Jax 0.6.0
This works great for Jax 0.5.3 giving me the speed from older versions, however with Jax 0.6.0 the script seems to crash when calling the ODE solver. Removing the flag allows my code to run, but noticeably slower than older jax version with the flag.
Have other people experienced this and have any suggestions?
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the reported diffrax ODE-solver crash with JAX 0.6.0 while setting XLA_FLAGS to --xla_cpu_use_thunk_runtime=false, then compare it with JAX 0.5.3 and with the flag removed. Done means determining whether the flag is incompatible and documenting a confirmed fix or workaround.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100