Glossary

As a reference, we’ll provide a glossary of some common technical terms used in the course.

JAX

Some common JAX terms are presented here. These terms may or may not match with the common terminologies.

• Asynchronous Dispatch: The phenomenon in which “jitted” function doesn’t await the complete calculation and passes it on to the next calculation.

• Device: A generic term for either CPU, GPU or TPU.

• DeviceArray : JAX’s analog of the numpy.ndarray.

• jaxpr: JAX Expressions (or jaxpr) are the intermediate representations of a computation graph.

• JIT: Just In Time compilation. It is performed in JAX using XLA.

• Pytrees: A tree-like structure built out of container-like Python objects.

• static: Compile-time computation in a JIT compilation and hence not traced.

• TPU: Tensor Processing Unit,

• Tracer: An object determining the sequence of operations performed by a Python function.

• Transformation: A ...