text: # Copyright 2018 Google LLC
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  # -*- coding: utf-8 -*-
  #
  # Configuration file for the Sphinx documentation builder.
  #
  # This file does only contain a selection of the most common options. For a
  # full list see the documentation:
  # http://www.sphinx-doc.org/en/master/config
  
  # -- Path setup --------------------------------------------------------------
  
  # If extensions (or modules to document with autodoc) are in another directory,
  # add these directories to sys.path here. If the directory is relative to the
  # documentation root, use os.path.abspath to make it absolute, like shown here.
  #
  import os
  import sys
  import typing
  
  sys.path.insert(0, os.path.abspath('..'))
  
  
  # Workaround to avoid expanding type aliases. See:
  # https://github.com/sphinx-doc/sphinx/issues/6518#issuecomment-589613836
  
  # When building docs, enable `from __future__ import annotations` everywhere.
  def _rewrite(p):
    with open(p) as f:
      contents = f.read()
    with open(p, 'w') as f:
      f.write('from __future__ import annotations\n')
      f.write(contents)
  
  if 'READTHEDOCS' in os.environ:
    for path, dirs, files in os.walk('../jax/'):
      for file in files:
        if file.endswith('.py'):
          _rewrite(os.path.abspath(os.path.join(path, file)))
  
  # Monkey patch for the typing module to prevent it from expanding type aliases.
  typing.get_type_hints = lambda obj, *unused: obj.__annotations__
  
  # -- Project information -----------------------------------------------------
  
  project = 'JAX'
  copyright = '2019, Google LLC. NumPy and SciPy documentation are copyright the respective authors.'
  author = 'The JAX authors'
  
  # The short X.Y version
  version = ''
  # The full version, including alpha/beta/rc tags
  release = ''
  
  
  # -- General configuration ---------------------------------------------------
  
  # If your documentation needs a minimal Sphinx version, state it here.
  #
  # needs_sphinx = '1.0'
  
  # Add any Sphinx extension module names here, as strings. They can be
  # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
  # ones.
  extensions = [
      'sphinx.ext.autodoc',
      'sphinx.ext.autosummary',
      'sphinx.ext.intersphinx',
      'sphinx.ext.mathjax',
      'sphinx.ext.napoleon',
      'sphinx.ext.viewcode',
      'nbsphinx',
      'sphinx_autodoc_typehints',
  ]
  
  intersphinx_mapping = {
      'python': ('https://docs.python.org/3/', None),
      'numpy': ('https://docs.scipy.org/doc/numpy/', None),
      'scipy': ('https://docs.scipy.org/doc/scipy/reference/', None),
  }
  
  # Add any paths that contain templates here, relative to this directory.
  templates_path = ['_templates']
  
  # The suffix(es) of source filenames.
  # You can specify multiple suffix as a list of string:
  #
  source_suffix = '.rst'
  
  # The master toctree document.
  master_doc = 'index'
  
  # The language for content autogenerated by Sphinx. Refer to documentation
  # for a list of supported languages.
  #
  # This is also used if you do content translation via gettext catalogs.
  # Usually you set "language" from the command line for these cases.
  language = None
  
  # List of patterns, relative to source directory, that match files and
  # directories to ignore when looking for source files.
  # This pattern also affects html_static_path and html_extra_path.
  exclude_patterns = [
      # Slow notebook: long time to load tf.ds
      'notebooks/neural_network_with_tfds_data.ipynb',
      # Slow notebook
      'notebooks/Neural_Network_and_Data_Loading.ipynb',
      'notebooks/score_matching.ipynb',
      'notebooks/maml.ipynb',
      # Fails with shape error in XL
      'notebooks/XLA_in_Python.ipynb',
      # Sometimes sphinx reads its own outputs as inputs!
      'build/html',
  ]
  
  # The name of the Pygments (syntax highlighting) style to use.
  pygments_style = None
  
  
  autosummary_generate = True
  napolean_use_rtype = False
  
  # -- Options for nbsphinx -----------------------------------------------------
  
  # Execute notebooks before conversion: 'always', 'never', 'auto' (default)
  # We execute all notebooks, exclude the slow ones using 'exclude_patterns'
  nbsphinx_execute = 'always'
  
  # Use this kernel instead of the one stored in the notebook metadata:
  #nbsphinx_kernel_name = 'python3'
  
  # List of arguments to be passed to the kernel that executes the notebooks:
  # nbsphinx_execute_arguments = []
  
  # If True, the build process is continued even if an exception occurs:
  #nbsphinx_allow_errors = True
  
  
  # Controls when a cell will time out (defaults to 30; use -1 for no timeout):
  nbsphinx_timeout = 180
  
  # Default Pygments lexer for syntax highlighting in code cells:
  #nbsphinx_codecell_lexer = 'ipython3'
  
  # Width of input/output prompts used in CSS:
  #nbsphinx_prompt_width = '8ex'
  
  # If window is narrower than this, input/output prompts are on separate lines:
  #nbsphinx_responsive_width = '700px'
  
  # This is processed by Jinja2 and inserted before each notebook
  nbsphinx_prolog = r"""
  {% set docname = 'docs/' + env.doc2path(env.docname, base=None) %}
  
  .. only:: html
  
      .. role:: raw-html(raw)
          :format: html
  
      .. nbinfo::
  
          Interactive online version:
          :raw-html:`<a href="https://colab.research.google.com/github/google/jax/blob/master/{{ docname }}"><img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg" style="vertical-align:text-bottom"></a>`
  
      __ https://github.com/google/jax/blob/
          {{ env.config.release }}/{{ docname }}
  """
  
  # This is processed by Jinja2 and inserted after each notebook
  # nbsphinx_epilog = r"""
  # """
  
  # Input prompt for code cells. "%s" is replaced by the execution count.
  #nbsphinx_input_prompt = 'In [%s]:'
  
  # Output prompt for code cells. "%s" is replaced by the execution count.
  #nbsphinx_output_prompt = 'Out[%s]:'
  
  # Specify conversion functions for custom notebook formats:
  #import jupytext
  #nbsphinx_custom_formats = {
  #    '.Rmd': lambda s: jupytext.reads(s, '.Rmd'),
  #}
  
  # Link or path to require.js, set to empty string to disable
  #nbsphinx_requirejs_path = ''
  
  # Options for loading require.js
  #nbsphinx_requirejs_options = {'async': 'async'}
  
  # mathjax_config = {
  #     'TeX': {'equationNumbers': {'autoNumber': 'AMS', 'useLabelIds': True}},
  # }
  
  # Additional files needed for generating LaTeX/PDF output:
  # latex_additional_files = ['references.bib']
  
  # -- Options for HTML output -------------------------------------------------
  
  # The theme to use for HTML and HTML Help pages.  See the documentation for
  # a list of builtin themes.
  #
  html_theme = 'sphinx_rtd_theme'
  
  # Theme options are theme-specific and customize the look and feel of a theme
  # further.  For a list of options available for each theme, see the
  # documentation.
  html_theme_options = {
      'logo_only': True,
  }
  
  # The name of an image file (relative to this directory) to place at the top
  # of the sidebar.
  html_logo = '_static/jax_logo_250px.png'
  
  # Add any paths that contain custom static files (such as style sheets) here,
  # relative to this directory. They are copied after the builtin static files,
  # so a file named "default.css" will overwrite the builtin "default.css".
  html_static_path = ['_static']
  
  # Custom sidebar templates, must be a dictionary that maps document names
  # to template names.
  #
  # The default sidebars (for documents that don't match any pattern) are
  # defined by theme itself.  Builtin themes are using these templates by
  # default: ``['localtoc.html', 'relations.html', 'sourcelink.html',
  # 'searchbox.html']``.
  #
  # html_sidebars = {}
  
  
  # -- Options for HTMLHelp output ---------------------------------------------
  
  # Output file base name for HTML help builder.
  htmlhelp_basename = 'JAXdoc'
  
  
  # -- Options for LaTeX output ------------------------------------------------
  
  latex_elements = {
      # The paper size ('letterpaper' or 'a4paper').
      #
      # 'papersize': 'letterpaper',
  
      # The font size ('10pt', '11pt' or '12pt').
      #
      # 'pointsize': '10pt',
  
      # Additional stuff for the LaTeX preamble.
      #
      # 'preamble': '',
  
      # Latex figure (float) alignment
      #
      # 'figure_align': 'htbp',
  }
  
  # Grouping the document tree into LaTeX files. List of tuples
  # (source start file, target name, title,
  #  author, documentclass [howto, manual, or own class]).
  latex_documents = [
      (master_doc, 'JAX.tex', 'JAX Documentation',
       'The JAX authors', 'manual'),
  ]
  
  
  # -- Options for manual page output ------------------------------------------
  
  # One entry per manual page. List of tuples
  # (source start file, name, description, authors, manual section).
  man_pages = [
      (master_doc, 'jax', 'JAX Documentation',
       [author], 1)
  ]
  
  
  # -- Options for Texinfo output ----------------------------------------------
  
  # Grouping the document tree into Texinfo files. List of tuples
  # (source start file, target name, title, author,
  #  dir menu entry, description, category)
  texinfo_documents = [
      (master_doc, 'JAX', 'JAX Documentation',
       author, 'JAX', 'One line description of project.',
       'Miscellaneous'),
  ]
  
  
  # -- Options for Epub output -------------------------------------------------
  
  # Bibliographic Dublin Core info.
  epub_title = project
  
  # The unique identifier of the text. This can be a ISBN number
  # or the project homepage.
  #
  # epub_identifier = ''
  
  # A unique identification for the text.
  #
  # epub_uid = ''
  
  # A list of files that should not be packed into the epub file.
  epub_exclude_files = ['search.html']
  
  
  # -- Extension configuration -------------------------------------------------
  
  # Tell sphinx-autodoc-typehints to generate stub parameter annotations including
  # types, even if the parameters aren't explicitly documented.
  always_document_param_types = True
  