Skip to main content

Objax is a machine learning framework that provides an Object Oriented layer for JAX.

Project description

Objax

Tutorials | Install | Documentation | Philosophy

This is not an officially supported Google product.

Objax is an open source machine learning framework that accelerates research and learning thanks to a minimalist object-oriented design and a readable code base. Its name comes from the contraction of Object and JAX -- a popular high-performance framework. Objax is designed by researchers for researchers with a focus on simplicity and understandability. Its users should be able to easily read, understand, extend, and modify it to fit their needs.

This is the developer repository of Objax, there is very little user documentation here, for the full documentation go to objax.readthedocs.io.

You can find READMEs in the subdirectory of this project, for example:

User installation guide

You install Objax using pip as follows:

pip install --upgrade objax

Objax supports GPUs but assumes that you already have some version of CUDA installed. Here are the extra steps required to install CUDA-enabled jaxlib (jaxlib releases require CUDA 11.2 or newer):

RELEASE_URL="https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
JAX_VERSION=`python3 -c 'import jax; print(jax.__version__)'`
pip uninstall -y jaxlib
pip install -f $RELEASE_URL jax[cuda]==$JAX_VERSION

For more installation options, see https://github.com/google/jax#pip-installation-gpu-cuda

Useful environment configurations

Here are a few useful options:

# Prevent JAX from taking the whole GPU memory
# (useful if you want to run several programs on a single GPU)
export XLA_PYTHON_CLIENT_PREALLOCATE=false

Testing your installation

You can test your installation by running the code below:

import jax
import objax

print(f'Number of GPUs {jax.device_count()}')

x = objax.random.normal(shape=(100, 4))
m = objax.nn.Linear(nin=4, nout=5)
print('Matrix product shape', m(x).shape)  # (100, 5)

x = objax.random.normal(shape=(100, 3, 32, 32))
m = objax.nn.Conv2D(nin=3, nout=4, k=3)
print('Conv2D return shape', m(x).shape)  # (100, 4, 32, 32)

Typically if you get errors running this using CUDA, it probably means your installation of CUDA or CuDNN has issues.

Runing code examples

Clone the code repository:

git clone https://github.com/google/objax.git
cd objax/examples

Citing Objax

To cite this repository:

@software{objax2020github,
  author = {{Objax Developers}},
  title = {{Objax}},
  url = {https://github.com/google/objax},
  version = {1.2.0},
  year = {2020},
}

Developer documentation

Here is information about development setup and a guide on adding new code.

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

objax-1.8.0.tar.gz (59.9 kB view details)

Uploaded Source

Built Distribution

objax-1.8.0-py3-none-any.whl (86.7 kB view details)

Uploaded Python 3

File details

Details for the file objax-1.8.0.tar.gz.

File metadata

  • Download URL: objax-1.8.0.tar.gz
  • Upload date:
  • Size: 59.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.11.5

File hashes

Hashes for objax-1.8.0.tar.gz
Algorithm Hash digest
SHA256 55d34f393ac9e57f9ff73be4d7c2d9efdb0ba8fb2c27d7564f2bd06fa19e90da
MD5 e845d1e4f1daa350483315248768a4f8
BLAKE2b-256 200ccbadc1a538f5ad3baa529b9acb75f5433d291a5bc23787ed239bb5b24d29

See more details on using hashes here.

File details

Details for the file objax-1.8.0-py3-none-any.whl.

File metadata

  • Download URL: objax-1.8.0-py3-none-any.whl
  • Upload date:
  • Size: 86.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.11.5

File hashes

Hashes for objax-1.8.0-py3-none-any.whl
Algorithm Hash digest
SHA256 7b0e3b28751529adc1402a05cc1e71f6e01b1476352d44cbdc4094d2586d1d1c
MD5 c6fc0d1983cfe3b0fe41d9cada53bb18
BLAKE2b-256 d4664847bec2b72c61d4d131e0dbff9ecf1b9eced084441e0cdb8a4f162738b3

See more details on using hashes here.

Supported by

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