Installing OJAX

Dependencies

OJAX is written in pure Python and requires Python 3.10+ (same as JAX). typing-extensions package is also loaded for using the newer Python typing features to properly do type hinting and mypy checks.

Apart from this, OJAX only depends on JAX, which can be installed following the JAX official installation guide.

Since JAX is still constantly changing the API, OJAX releases will align with the API of the concurrent latest stable JAX releases.

Installing OJAX

Once JAX is installed, OJAX can be easily installed with pip:

pip install ojax