Farama-Foundation / Farama-Foundation/Jumpy

Jax and Numpy call priority

Open
#23 7 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.