JAX Glossary of Terms
=====================

.. glossary::

    CPU
      Short for *Central Processing Unit*, CPUs are the standard computational architecture
      available in most computers. JAX can run computations on CPUs, but often can achieve
      much better performance on :term:`GPU` and :term:`TPU`.

    Device
      The generic name used to refer to the :term:`CPU`, :term:`GPU`, or :term:`TPU` used
      by JAX to perform computations.

    DeviceArray
      JAX's analog of the :class:`numpy.ndarray`. See :class:`jaxlib.xla_extension.DeviceArray`.

    forward-mode autodiff
      See :term:`JVP`

    functional programming
      A programming paradigm in which programs are defined by applying and composing
      :term:`pure functions<pure function>`. JAX is designed for use with functional programs.

    GPU
      Short for *Graphical Processing Unit*, GPUs were originally specialized for operations
      related to rendering of images on screen, but now are much more general-purpose. JAX is
      able to target GPUs for fast operations on arrays (see also :term:`CPU` and :term:`TPU`).

    jaxpr
      Short for *JAX Expression*, a jaxpr is an intermediate representation of a computation that
      is generated by JAX, and is forwarded to :term:`XLA` for compilation and execution.
      See :ref:`understanding-jaxprs` for more information. 

    JIT
      Short for *Just In Time* compilation, JIT in JAX generally refers to the compilation of
      array operations to :term:`XLA`, most often accomplished using :func:`jax.jit`.

    JVP
      Short for *Jacobian Vector Product*, also sometimes known as *forward-mode* automatic
      differentiation. For more details, see :ref:`jacobian-vector-product`. In JAX, JVP is
      a :term:`transformation` that is implemented via :func:`jax.jvp`. See also :term:`VJP`.

    pure function
      A pure function is a function whose outputs are based only on its inputs, and which has
      no side-effects. JAX's :term:`transformation` model is designed to work with pure functions.
      See also :term:`functional programming`.

    reverse-mode autodiff
      See :term:`VJP`.

    SPMD
      Short for *Single Program Multi Data*, it refers to a parallel computation technique in which
      the same computation (e.g., the forward pass of a neural net) is run on different input data
      (e.g., different inputs in a batch) in parallel on different devices (e.g., several TPUs).
      :func:`jax.pmap` is a JAX :term:`transformation` that implements SPMD parallelism.

    static
      In a :term:`JIT` compilation, a value that is not traced (see :term:`Tracer`). Also
      sometimes refers to compile-time computations on static values.
    
    TPU
      Short for *Tensor Processing Unit*, TPUs are chips specifically engineered for fast operations
      on N-dimensional tensors used in deep learning applications. JAX is able to target TPUs for
      fast operations on arrays (see also :term:`CPU` and :term:`GPU`).

    Tracer
      An object used as a standin for a JAX :term:`DeviceArray` in order to determine the
      sequence of operations performed by a Python function. Internally, JAX implements this
      via the :class:`jax.core.Tracer` class.

    transformation
      A higher-order function: that is, a function that takes a function as input and outputs
      a transformed function. Examples in JAX include :func:`jax.jit`, :func:`jax.vmap`, and
      :func:`jax.grad`.

    VJP
      Short for *Vector Jacobian Product*, also sometimes known as *reverse-mode* automatic
      differentiation. For more details, see :ref:`vector-jacobian-product`. In JAX, VJP is
      a :term:`transformation` that is implemented via :func:`jax.vjp`. See also :term:`JVP`.

    XLA
      Short for *Accelerated Linear Algebra*, XLA is a domain-specific compiler for linear
      algebra operations that is the primary backend for :term:`JIT`-compiled JAX code.
      See https://www.tensorflow.org/xla/.

    weak type
      A JAX data type that has the same type promotion semantics as Python scalars;
      see :ref:`weak-types`.