{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-1 fw=all -->\n",
    "\n",
    "# Documentation\n",
    "\n",
    "No matter how much of a framework's API we cover here,\n",
    "there will always be functions, classes, and arguments\n",
    "we never reach, and the libraries keep changing under us.\n",
    "So rather than try to memorize the API,\n",
    "the durable skill is getting good at *looking things up*:\n",
    "finding what exists, reading how it works,\n",
    "and confirming that it does what you think.\n",
    "This short section lays out a small, repeatable loop for exactly that,\n",
    "using tools built into Python and your notebook.\n",
    "\n",
    "The official documentation is always the source of truth;\n",
    "bookmark the reference and tutorial pages for the framework you use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-2 fw=jax -->\n",
    "\n",
    "For JAX these are the\n",
    "[API reference](https://jax.readthedocs.io/en/latest/) and the\n",
    "[tutorials](https://jax.readthedocs.io/en/latest/tutorials.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-3 fw=all -->\n",
    "\n",
    "For most day-to-day questions, though, you do not need to leave your\n",
    "notebook. Four moves, repeated until the call behaves, cover almost\n",
    "everything.\n",
    "\n",
    "![Four moves, repeated until the call does what you want: **discover** the names that exist, **inspect** a candidate's signature, **read** the docs or source when you need the *why*, and **verify** with a quick run.](../img/lookup-api-discovery-loop.svg)\n",
    "\n",
    "The examples below start from the standard import:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "lookup-api-documentation"
   },
   "source": [
    "import jax"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-4 fw=all -->\n",
    "\n",
    "## Discovering What Exists: `dir`\n",
    "\n",
    "When you know roughly *where* a tool should live but not what it is called,\n",
    "the `dir` function lists everything defined in a module.\n",
    "For instance, to see what is on offer for random sampling\n",
    "(we print the first twenty names):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "lookup-api-functions-and-classes-in-a-module"
   },
   "source": [
    "print([name for name in dir(jax.random) if not name.startswith('_')][:20])"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "['PRNGKey', 'ball', 'bernoulli', 'beta', 'binomial', 'bits', 'categorical', 'cauchy', 'chisquare', 'choice', 'clone', 'dirichlet', 'double_sided_maxwell', 'exponential', 'f', 'fold_in', 'gamma', 'generalized_normal', 'geometric', 'gumbel']\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-5 fw=all -->\n",
    "\n",
    "We can usually ignore names that begin and end with `__`\n",
    "(Python's special objects) or that start with a single `_`\n",
    "(internal helpers). The remaining names already hint\n",
    "at what the module offers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-6 fw=jax -->\n",
    "\n",
    "Here the names are samplers (`bernoulli`, `beta`, `cauchy`,\n",
    "`exponential`, `gamma`, ...) plus `PRNGKey`, which creates the explicit\n",
    "random key that every JAX sampler takes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-7 fw=all -->\n",
    "\n",
    "In a notebook you can get the same list interactively, filtered as you\n",
    "type, by writing the module name followed by a dot and pressing `Tab`;\n",
    "this is usually the fastest way to turn up a name.\n",
    "\n",
    "## Reading the Signature: `help`, `?`, and `??`\n",
    "\n",
    "Once you have a name, `help` prints its docstring:\n",
    "the arguments it takes, their defaults, what it returns,\n",
    "and often a short example. Let us look up the `ones` function,\n",
    "which we have used to build tensors:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "lookup-api-specific-functions-and-classes-1"
   },
   "source": [
    "help(jax.numpy.ones)"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Help on function ones in module jax.numpy:\n\nones(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\n    Create an array full of ones.\n\n    JAX implementation of :func:`numpy.ones`.\n\n    Args:\n      shape: int or sequence of ints specifying the shape of the created array.\n      dtype: optional dtype for the created array; defaults to float32 or float64\n        depending on the X64 configuration (see :ref:`default-dtypes`).\n      device: (optional) :class:`~jax.Device` or :class:`~jax.sharding.Sharding`\n        to which the created array will be committed. This argument exists for\n        compatibility with the :ref:`python-array-api`.\n      out_sharding: (optional) :class:`~jax.sharding.PartitionSpec` or :class:`~jax.NamedSharding`\n        representing the sharding of the created array (see `explicit sharding`_ for more details).\n        This argument exists for consistency with other array creation routines across JAX.\n        Specifying both ``out_sharding`` and ``device`` will result in an error.\n\n    Returns:\n      Array of the specified shape and dtype, with the given device/sharding if specified.\n\n    See also:\n      - :func:`jax.numpy.ones_like`\n      - :func:`jax.numpy.empty`\n      - :func:`jax.numpy.zeros`\n      - :func:`jax.numpy.full`\n\n    Examples:\n      >>> jnp.ones(4)\n      Array([1., 1., 1., 1.], dtype=float32)\n      >>> jnp.ones((2, 3), dtype=bool)\n      Array([[ True,  True,  True],\n             [ True,  True,  True]], dtype=bool)\n\n    .. _explicit sharding: https://docs.jax.dev/en/latest/parallel.html\n\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-8 fw=all -->\n",
    "\n",
    "The docstring tells us that `ones` creates a new tensor of the requested\n",
    "shape with every element set to 1.\n",
    "In a Jupyter notebook, two shortcuts make this quicker still:\n",
    "`ones?` opens the same docstring in a side pane,\n",
    "and `ones??` additionally displays the function's *source code*.\n",
    "The source is the final word when a docstring is terse or ambiguous,\n",
    "and reading it is one of the better ways to pick up idioms\n",
    "from high-quality libraries.\n",
    "\n",
    "## Verifying With a Quick Run\n",
    "\n",
    "Docstrings can be terse, and they occasionally drift out of date.\n",
    "The fastest way to be certain is to run a tiny example\n",
    "and look at the result:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "lookup-api-specific-functions-and-classes-2"
   },
   "source": [
    "jax.numpy.ones(4)"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([1., 1., 1., 1.], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-9 fw=all -->\n",
    "\n",
    "The shape and values are exactly what the docstring promised.\n",
    "Making this `discover → inspect → read → verify` loop a habit\n",
    "will carry you through the unfamiliar corners of any library,\n",
    "long after the specific functions in this book have changed.\n",
    "\n",
    "Coding assistants are often the quickest route to a first answer:\n",
    "ask \"how do I sample from a normal distribution in this framework?\"\n",
    "and you will usually get a function and a working call in seconds.\n",
    "Treat the suggestion the way you would a knowledgeable colleague's tip:\n",
    "a good starting point that still goes through the loop above.\n",
    "Glance at the signature with `help` or `?`, run a small example,\n",
    "and rely on the suggestion once it survives both.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Use `dir` on your framework's random-number module to find the routine\n",
    "   that samples from a *uniform* distribution. Read its signature with `help`\n",
    "   (or `?`), then call it to draw a $3 \\times 3$ tensor and confirm the values\n",
    "   lie in $[0, 1)$.\n",
    "1. You want to reduce a tensor along a single axis but cannot remember the\n",
    "   keyword. Look up your framework's `sum` (or `reduce_sum`) with `help`,\n",
    "   identify the argument that selects the axis, and verify on a $2 \\times 3$\n",
    "   tensor that summing over each axis gives the shape you predicted.\n",
    "1. Ask a coding assistant \"how do I concatenate two tensors along a new axis\n",
    "   in my framework?\" Then run its answer through the\n",
    "   discover&nbsp;&rarr;&nbsp;inspect&nbsp;&rarr;&nbsp;read&nbsp;&rarr;&nbsp;verify\n",
    "   loop: does the suggested function exist (`dir`), does its signature match\n",
    "   the claim (`help`/`?`), and does a tiny example do what you expect?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=lookup-api-md-10 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/17972)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
