Skip to main content

Common backend for JAX or numpy.

Project description

Jumpy is a common backend for JAX or NumPy:

  • A Jumpy function returns a JAX outputs if given a JAX inputs
  • A Jumpy function returns a JAX outputs if jitted
  • Otherwise a jumpy function returns NumPy outputs

Jumpy lets you write framework agnostic code that is easy to debug by running as raw Numpy, but is just as performant as JAX when jitted.

We maintain this repository primarily so to enable writing Gymnasium and PettingZoo wrappers that can be applied to both standard NumPy or hardware accelerated Jax based environments, however this package can be used for many more things.

Installing Jumpy

To install Jumpy from pypi:

python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install brax-jumpy

Alternatively, to install Jumpy from source, clone this repo, cd to it, and then:

python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jax_jumpy-0.2.0.tar.gz (94.1 kB view details)

Uploaded Source

Built Distribution

jax_jumpy-0.2.0-py3-none-any.whl (11.1 kB view details)

Uploaded Python 3

File details

Details for the file jax_jumpy-0.2.0.tar.gz.

File metadata

  • Download URL: jax_jumpy-0.2.0.tar.gz
  • Upload date:
  • Size: 94.1 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.9.13

File hashes

Hashes for jax_jumpy-0.2.0.tar.gz
Algorithm Hash digest
SHA256 8ffe4e7609461d13decc97269ff48a3ee8d7570c96cbcc82487f893a3dd65093
MD5 cab34afe9fb966d3fd45a4a2d7a97f82
BLAKE2b-256 7f10f8530524cd7d6b076ca1e99ae6082d545e9304332ce85aeb78765e7c2a5f

See more details on using hashes here.

File details

Details for the file jax_jumpy-0.2.0-py3-none-any.whl.

File metadata

  • Download URL: jax_jumpy-0.2.0-py3-none-any.whl
  • Upload date:
  • Size: 11.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.9.13

File hashes

Hashes for jax_jumpy-0.2.0-py3-none-any.whl
Algorithm Hash digest
SHA256 eaa7587e3fda30278011ca4e815bf9ac6b71db29502c0fa15908eb532c206a9e
MD5 386f75e3debb11a07d972ce26e5d67eb
BLAKE2b-256 54313467f6dfeac9c31b782b089b18e78b5f3b01aed6e1e1b5601745b1a3e6ae

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page