Farama-Foundation / Farama-Foundation/Jumpy
Jax and Numpy call priority
- Dominant language
- Python
- Stars
- 56
- Forks
- 9
- PR merge metrics
- No merged PRs in 30d
Description
Hi,
I was wondering whether it was possible to force the use of numpy even though jax was installed. Namely, I would like to run some parts of my code with numpy, while other parts should still be run with the jax backend.
[This line](https://github.com/Farama-Foundation/Jumpy/blob/main/jumpy.py#L24) seems to suggest that it globally determines what backend is used. I could change the `_has_jax` variable before and after, but it seems a hacky.
If there is a better way, let me know!
A suggestion could be to use a context manager for this, similar to how `chex.fake_jit` works:
```
import jumpy as jnp
with jnp.use_jax():
....
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.