User Guide

By default, JAX does all floating computations in 32-bit. We should configure it to use 64-bit at the beginning of any program.

# Configure JAX for 64-bit computing
from jax.config import config
config.update("jax_enable_x64", True)

NumPy features from JAX

import jax.numpy as jnp

We often need random keys generator from JAX. It is convenient to generate a set of keys in advance in a script and use them as needed.

from jax import random
# Some keys for generating random numbers
key = random.PRNGKey(0)
keys = random.split(key, 4)

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Import Conventions

CR-Suite contains a large collection of functions organized the form of modules. Here we summarize basic conventions of how to load individual modules. These conventions will be followed in the coding examples in rest of the book.

CR-Nimble

Top level functions:

import cr.nimble as crn

Digital Signal Processing functions:

import cr.nimble.dsp as crdsp

CR-Wavelets

import cr.wavelets as crwt

CR-Sparse

Top level functions:

import cr.sparse as crs

Linear operators:

import cr.sparse.lop as crlop