Fixes bugs, adds new functions, moves dtypes to `jumpy.numpy` package

Farama-Foundation / Jumpy

On-the-fly conversions between Jax and NumPy tensors

Apache License 2.0

47 stars 9 forks source link

Closed bheijden closed 1 year ago

bheijden commented 1 year ago

This PR addresses the following:

bug fixes for jumpy.lax.fori_loop and jumpy.lax.cond.
Moves all dtypes from jumpy to jumpy.numpy, following the module structure of jax.numpy and numpy.
Added some new implementations to jumpy.lax module (e.g.stop_gradient and switch).
Fixes non-deterministic behavior with jitted vs non-jitted functions. Closes #29.