1.7  Documentation

A framework API contains more functions, classes, and arguments than any single text can cover, and it changes across releases. Effective use therefore depends on being able to find an operation, inspect its interface, read its documentation, and verify its behavior with a small example. This section demonstrates that procedure using Python and notebook tools.

The official reference and tutorial pages document the supported interface for each framework.

For PyTorch these are the API reference and the tutorials.

For TensorFlow these are the API reference and the tutorials.

For JAX these are the API reference and the tutorials.

For MXNet these are the API reference and the tutorials. One caveat: these pages document version 1.9.1 (the last release with hosted documentation), which foregrounds the legacy mx.nd interface, whereas this book uses the NumPy-style np/npx API of MXNet 2. For MXNet specifics, the in-notebook loop below is often the more reliable reference.

Many routine questions can be answered without leaving the notebook. The procedure has four steps: discover, inspect, read, and verify.

Figure 1.7.1: A four-step procedure for using an unfamiliar API: discover available names, inspect a candidate’s signature, read its documentation or source, and verify its behavior with a small example.

The examples below start from the standard import:

import torch
from pprint import pprint
import tensorflow as tf
from pprint import pprint
import jax
from pprint import pprint
from mxnet import np
from pprint import pprint

1.7.1 Discovering What Exists: dir

When you know roughly where a tool should live but not what it is called, the dir function lists the names defined in a module. For instance, to inspect the available random-sampling tools (we print the first twenty names):

pprint([name for name in dir(torch.distributions)
        if not name.startswith('_')][:20], compact=True)
['AbsTransform', 'AffineTransform', 'Bernoulli', 'Beta', 'Binomial',
 'CatTransform', 'Categorical', 'Cauchy', 'Chi2', 'ComposeTransform',
 'ContinuousBernoulli', 'CorrCholeskyTransform',
 'CumulativeDistributionTransform', 'Dirichlet', 'Distribution', 'ExpTransform',
 'Exponential', 'ExponentialFamily', 'FisherSnedecor', 'Gamma']
pprint([name for name in dir(tf.random)
        if not name.startswith('_')][:20], compact=True)
['Algorithm', 'Generator', 'all_candidate_sampler', 'categorical',
 'create_rng_state', 'experimental', 'fixed_unigram_candidate_sampler',
 'fold_in', 'gamma', 'get_global_generator',
 'learned_unigram_candidate_sampler', 'log_uniform_candidate_sampler', 'normal',
 'poisson', 'set_global_generator', 'set_seed', 'shuffle', 'split',
 'stateless_binomial', 'stateless_categorical']
pprint([name for name in dir(jax.random)
        if not name.startswith('_')][:20], compact=True)
['PRNGKey', 'ball', 'bernoulli', 'beta', 'binomial', 'bits', 'categorical',
 'cauchy', 'chisquare', 'choice', 'clone', 'dirichlet', 'double_sided_maxwell',
 'exponential', 'f', 'fold_in', 'gamma', 'generalized_normal', 'geometric',
 'gumbel']
pprint([name for name in dir(np.random)
        if not name.startswith('_')][:20], compact=True)
['beta', 'chisquare', 'choice', 'exponential', 'f', 'gamma', 'gumbel',
 'integer_types', 'laplace', 'logistic', 'lognormal', 'multinomial',
 'multivariate_normal', 'normal', 'pareto', 'power', 'rand', 'randint', 'randn',
 'rayleigh']

We can usually ignore names that begin and end with __ (Python’s special objects) or that start with a single _ (internal helpers). The remaining names already hint at what the module offers.

Here the names are distribution classes such as Bernoulli, Categorical, and Gamma; each can be instantiated and then sampled from. The Transform entries build new distributions by transforming existing ones.

Here we can spot samplers such as gamma, normal, and poisson, next to utilities like Generator and set_seed that manage the random state.

Here the names are samplers (bernoulli, beta, cauchy, exponential, gamma, …) plus PRNGKey, which creates the explicit random key that every JAX sampler takes.

Here almost every name is a sampler: draws from classical distributions (beta, gamma, multinomial, normal, …) alongside NumPy-style conveniences such as rand, randint, and randn.

In a notebook you can get the same list interactively, filtered as you type, by writing the module name followed by a dot and pressing Tab; this often locates a name quickly.

1.7.2 Reading the Signature: help, ?, and ??

Once you have a name, help prints its docstring: the arguments it takes, their defaults, what it returns, and often a short example. Let us look up the ones function, which we have used to build tensors:

help(torch.ones)
Help on built-in function ones in module torch:

ones(...)
    ones(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor

    Returns a tensor filled with the scalar value `1`, with the shape defined
    by the variable argument :attr:`size`.

    Args:
        size (int...): a sequence of integers defining the shape of the output tensor.
            Can be a variable number of arguments or a collection like a list or tuple.

    Keyword arguments:
        out (Tensor, optional): the output tensor.
        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.
            Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`).
        layout (:class:`torch.layout`, optional): the desired layout of returned Tensor.
            Default: ``torch.strided``.
        device (:class:`torch.device`, optional): the desired device of returned tensor.
            Default: if ``None``, uses the current device for the default tensor type
            (see :func:`torch.set_default_device`). :attr:`device` will be the CPU
            for CPU tensor types and the current CUDA device for CUDA tensor types.
        requires_grad (bool, optional): If autograd should record operations on the
            returned tensor. Default: ``False``.

    Example::

        >>> torch.ones(2, 3)
        tensor([[ 1.,  1.,  1.],
                [ 1.,  1.,  1.]])

        >>> torch.ones(5)
        tensor([ 1.,  1.,  1.,  1.,  1.])
help(tf.ones)
Help on function ones in module tensorflow.python.ops.array_ops:

ones(shape, dtype=tf.float32, name=None, layout=None)
    Creates a tensor with all elements set to one (1).

    See also `tf.ones_like`, `tf.zeros`, `tf.fill`, `tf.eye`.

    This operation returns a tensor of type `dtype` with shape `shape` and
    all elements set to one.

    >>> tf.ones([3, 4], tf.int32)
    <tf.Tensor: shape=(3, 4), dtype=int32, numpy=
    array([[1, 1, 1, 1],
           [1, 1, 1, 1],
           [1, 1, 1, 1]], dtype=int32)>

    Args:
      shape: A `list` of integers, a `tuple` of integers, or a 1-D `Tensor` of
        type `int32`.
      dtype: Optional DType of an element in the resulting `Tensor`. Default is
        `tf.float32`.
      name: Optional string. A name for the operation.
      layout: Optional, `tf.experimental.dtensor.Layout`. If provided, the result
        is a [DTensor](https://www.tensorflow.org/guide/dtensor_overview) with the
        provided layout.

    Returns:
      A `Tensor` with all elements set to one (1).
help(jax.numpy.ones)
Help on function ones in module jax.numpy:

ones(shape: Any, dtype: Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, NoneType] = None, *, device: jaxlib._jax.Device | jaxlib._jax.Sharding | None = None, out_sharding: jax.sharding.NamedSharding | jax.P | None = None) -> jax.Array
    Create an array full of ones.

    JAX implementation of :func:`numpy.ones`.

    Args:
      shape: int or sequence of ints specifying the shape of the created array.
      dtype: optional dtype for the created array; defaults to float32 or float64
        depending on the X64 configuration (see :ref:`default-dtypes`).
      device: (optional) :class:`~jax.Device` or :class:`~jax.sharding.Sharding`
        to which the created array will be committed. This argument exists for
        compatibility with the :ref:`python-array-api`.
      out_sharding: (optional) :class:`~jax.sharding.PartitionSpec` or :class:`~jax.NamedSharding`
        representing the sharding of the created array (see `explicit sharding`_ for more details).
        This argument exists for consistency with other array creation routines across JAX.
        Specifying both ``out_sharding`` and ``device`` will result in an error.

    Returns:
      Array of the specified shape and dtype, with the given device/sharding if specified.

    See also:
      - :func:`jax.numpy.ones_like`
      - :func:`jax.numpy.empty`
      - :func:`jax.numpy.zeros`
      - :func:`jax.numpy.full`

    Examples:
      >>> jnp.ones(4)
      Array([1., 1., 1., 1.], dtype=float32)
      >>> jnp.ones((2, 3), dtype=bool)
      Array([[ True,  True,  True],
             [ True,  True,  True]], dtype=bool)

    .. _explicit sharding: https://docs.jax.dev/en/latest/parallel.html
help(np.ones)
Help on function ones in module mxnet.numpy:

ones(shape, dtype=None, order='C', device=None)
    Return a new array of given shape and type, filled with ones.
    This function currently only supports storing multi-dimensional data
    in row-major (C-style).

    Parameters
    ----------
    shape : int or tuple of int
        The shape of the empty array.
    dtype : str or numpy.dtype, optional
        An optional value type. Default is depend on your current default dtype.
        When npx.is_np_default_dtype() returns False, default dtype is float32;
        When npx.is_np_default_dtype() returns True, default dtype is float64.
        Note that this behavior is different from NumPy's `ones` function where
        `float64` is the default value.
    order : {'C'}, optional, default: 'C'
        How to store multi-dimensional data in memory, currently only row-major
        (C-style) is supported.
    device : Device, optional
        Device context on which the memory is allocated. Default is
        `mxnet.device.current_device()`.

    Returns
    -------
    out : ndarray
        Array of ones with the given shape, dtype, and device.

    Examples
    --------
    >>> np.ones(5)
    array([1., 1., 1., 1., 1.])

    >>> np.ones((5,), dtype=int)
    array([1, 1, 1, 1, 1], dtype=int64)

    >>> np.ones((2, 1))
    array([[1.],
           [1.]])

    >>> s = (2,2)
    >>> np.ones(s)
    array([[1., 1.],
           [1., 1.]])

The docstring tells us that ones creates a new tensor of the requested shape with every element set to 1. In a Jupyter notebook, two shortcuts make this quicker still: ones? opens the same docstring in a side pane, and ones?? additionally displays the function’s source code. When a docstring is terse or ambiguous, the source can clarify the implementation. Reading library source also reveals established usage idioms.

1.7.3 Verifying With a Quick Run

Docstrings can be terse, and they occasionally drift out of date. A small example can verify the behavior directly:

torch.ones(4)
tensor([1., 1., 1., 1.])
tf.ones(4)
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([1., 1., 1., 1.], dtype=float32)>
jax.numpy.ones(4)
Array([1., 1., 1., 1.], dtype=float32)
np.ones(4)
array([1., 1., 1., 1.])

The result has the documented shape and values. Making this discover → inspect → read → verify loop a habit will carry you through the unfamiliar corners of any library, long after the specific functions in this book have changed.

Coding assistants can propose an API call for a question such as “how do I sample from a normal distribution in this framework?” Treat the result as an unverified candidate that still goes through the loop above. Glance at the signature with help or ?, run a small example, and accept the suggestion only once it survives both.

1.7.4 Summary

For an unfamiliar operation, discover candidate names, inspect the signature, read the relevant documentation, and verify the behavior with the smallest run that exposes shapes, dtypes, and values. This procedure remains valid as library APIs change and applies equally to suggestions from search or coding assistants.

1.7.5 Exercises

  1. [code] Scale versus rate. The gamma distribution with shape \(k = 2\) and rate \(\beta = 3\) has mean \(k/\beta = 2/3\) and variance \(k/\beta^2 = 2/9\). Use dir on torch.distributions to find its class and read the signature with help (or ?) to find out how the class is parameterized. Draw 10,000 samples with \(k = 2\) and \(\beta = 3\) and confirm the mean and variance from the samples.
  2. [code] Standard-deviation conventions. For \(n\) numbers whose squared deviations from their mean sum to \(S\), a library computes the standard deviation either as \(\sqrt{S/n}\) or as \(\sqrt{S/(n-1)}\). Section 1.2 standardized columns with the pandas method std. Compute the standard deviation of the values 1.0, 2.0, 3.0 and 4.0 with pd.Series.std, with torch.std, and with numpy.std. Use help on each function to find out how it sets the divisor, state the default divisor of each, and check the three results against a hand computation. Then make all three results agree.
  3. [code] Checking a coding assistant’s answer. Ask a coding assistant the two questions below. Check each function it names in three steps: dir or hasattr shows that the name exists, help or ? shows that the signature fits the suggested call, and a tiny example shows the claimed behavior.
    1. Ask how to split a PyTorch tensor into three equal parts along an axis.
    2. Ask for a single PyTorch function that reorders the axes of a tensor so that their lengths increase, an operation no library function provides. State which of the three checks, if any, each named function fails. Then build the operation from functions that pass all three checks and test it on a tensor of shape \((4, 2, 3)\).
  1. [code] Scale versus rate. The gamma distribution with shape \(k = 2\) and rate \(\beta = 3\) has mean \(k/\beta = 2/3\) and variance \(k/\beta^2 = 2/9\). Use dir on tf.random to find its sampler and read the signature with help (or ?) to find out how the sampler is parameterized. Draw 10,000 samples with \(k = 2\) and \(\beta = 3\) and confirm the mean and variance from the samples.
  2. [code] Standard-deviation conventions. For \(n\) numbers whose squared deviations from their mean sum to \(S\), a library computes the standard deviation either as \(\sqrt{S/n}\) or as \(\sqrt{S/(n-1)}\). Section 1.2 standardized columns with the pandas method std. Compute the standard deviation of the values 1.0, 2.0, 3.0 and 4.0 once with pd.Series.std and once with tf.math.reduce_std. Use help on both functions to find out how each one sets the divisor, state the default divisor of each, and check both results against a hand computation. Then make the two results agree.
  3. [code] Checking a coding assistant’s answer. Ask a coding assistant the two questions below. Check each function it names in three steps: dir or hasattr shows that the name exists, help or ? shows that the signature fits the suggested call, and a tiny example shows the claimed behavior.
    1. Ask how to split a TensorFlow tensor into three equal parts along an axis.
    2. Ask for a single TensorFlow function that reorders the axes of a tensor so that their lengths increase, an operation no library function provides. State which of the three checks, if any, each named function fails. Then build the operation from functions that pass all three checks and test it on a tensor of shape \((4, 2, 3)\).
  1. [code] Scale versus rate. The gamma distribution with shape \(k = 2\) and rate \(\beta = 3\) has mean \(k/\beta = 2/3\) and variance \(k/\beta^2 = 2/9\). Use dir on jax.random to find its sampler and read the signature with help (or ?) to find out how the sampler is parameterized. Draw 10,000 samples with \(k = 2\) and \(\beta = 3\), using an explicit key, and confirm the mean and variance from the samples.
  2. [code] Standard-deviation conventions. For \(n\) numbers whose squared deviations from their mean sum to \(S\), a library computes the standard deviation either as \(\sqrt{S/n}\) or as \(\sqrt{S/(n-1)}\). Section 1.2 standardized columns with the pandas method std. Compute the standard deviation of the values 1.0, 2.0, 3.0 and 4.0 once with pd.Series.std and once with jax.numpy.std. Use help on both functions to find out how each one sets the divisor, state the default divisor of each, and check both results against a hand computation. Then make the two results agree.
  3. [code] Checking a coding assistant’s answer. Ask a coding assistant the two questions below. Check each function it names in three steps: dir or hasattr shows that the name exists, help or ? shows that the signature fits the suggested call, and a tiny example shows the claimed behavior.
    1. Ask how to split a JAX array into three equal parts along an axis.
    2. Ask for a single JAX function that reorders the axes of an array so that their lengths increase, an operation no library function provides. State which of the three checks, if any, each named function fails. Then build the operation from functions that pass all three checks and test it on an array of shape \((4, 2, 3)\).
  1. [code] Scale versus rate. The gamma distribution with shape \(k = 2\) and rate \(\beta = 3\) has mean \(k/\beta = 2/3\) and variance \(k/\beta^2 = 2/9\). Use dir on np.random to find its sampler and read the signature with help (or ?) to find out how the sampler is parameterized. Draw 10,000 samples with \(k = 2\) and \(\beta = 3\) and confirm the mean and variance from the samples.
  2. [code] Standard-deviation conventions. For \(n\) numbers whose squared deviations from their mean sum to \(S\), a library computes the standard deviation either as \(\sqrt{S/n}\) or as \(\sqrt{S/(n-1)}\). Section 1.2 standardized columns with the pandas method std. Compute the standard deviation of the values 1.0, 2.0, 3.0 and 4.0 once with pd.Series.std and once with np.std on an MXNet array. Use help on both functions to find out how each one sets the divisor, state the default divisor of each, and check both results against a hand computation. Then make the two results agree.
  3. [code] Checking a coding assistant’s answer. Ask a coding assistant the two questions below. Check each function it names in three steps: dir or hasattr shows that the name exists, help or ? shows that the signature fits the suggested call, and a tiny example shows the claimed behavior.
    1. Ask how to split an MXNet tensor into three equal parts along an axis.
    2. Ask for a single MXNet function that reorders the axes of a tensor so that their lengths increase, an operation no library function provides. State which of the three checks, if any, each named function fails. Then build the operation from functions that pass all three checks and test it on a tensor of shape \((4, 2, 3)\).
  1. [code] Reading the source. The docstring of d2l.plot, defined in Section 1.4, describes the function only as “Plot data points.” After importing d2l as in that section, predict how many curves d2l.plot(x, Y) draws when x is a 1-D tensor of length 50 and Y is a list of five tensors of the same length. Read the source with d2l.plot?? or your editor’s go-to-definition, find the line that decides the number of curves, and confirm your reading with the call. Then pass an argument that makes all five curves appear.