mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 03:46:06 +00:00
Update doc landing page
Co-authored-by: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com> Co-authored-by: Jake VanderPlas <jakevdp@google.com>
This commit is contained in:
Ravin Kumar
parent
eb887f42c4
commit
2257e2075d
21
docs/_static/style.css
vendored
21
docs/_static/style.css
vendored
@ -1,5 +1,26 @@
|
||||
@import url("theme.css");
|
||||
|
||||
:root {
|
||||
--block-bg-opacity: .5;
|
||||
}
|
||||
|
||||
.wy-side-nav-search {
|
||||
background-color: #fff;
|
||||
}
|
||||
|
||||
.getting-started {
|
||||
background-color: rgba(78, 150, 253, var(--block-bg-opacity));
|
||||
}
|
||||
|
||||
.user-guides {
|
||||
background-color: rgba(0, 169, 154, var(--block-bg-opacity));
|
||||
}
|
||||
|
||||
.developer-docs {
|
||||
background-color: rgba(171, 0, 182, var(--block-bg-opacity));
|
||||
}
|
||||
|
||||
.key-ideas
|
||||
{
|
||||
border: 0px
|
||||
}
|
||||
20
docs/advanced_guide.rst
Normal file
20
docs/advanced_guide.rst
Normal file
@ -0,0 +1,20 @@
|
||||
.. _advanced_guide:
|
||||
|
||||
Advanced Tutorials
|
||||
==================
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
notebooks/autodiff_cookbook
|
||||
multi_process
|
||||
notebooks/Distributed_arrays_and_automatic_parallelization
|
||||
notebooks/vmapped_log_probs
|
||||
notebooks/neural_network_with_tfds_data
|
||||
notebooks/Custom_derivative_rules_for_Python_code
|
||||
notebooks/How_JAX_primitives_work
|
||||
notebooks/Writing_custom_interpreters_in_Jax
|
||||
notebooks/Neural_Network_and_Data_Loading
|
||||
notebooks/xmap_tutorial
|
||||
notebooks/external_callbacks
|
||||
Custom_Operation_for_GPUs
|
||||
47
docs/beginner_guide.rst
Normal file
47
docs/beginner_guide.rst
Normal file
@ -0,0 +1,47 @@
|
||||
:orphan:
|
||||
|
||||
.. _beginner-guide:
|
||||
|
||||
|
||||
Beginner Guide
|
||||
==============
|
||||
|
||||
Welcome to the beginners guide for JAX.
|
||||
On this page we will introduce you to the key ideas of JAX,
|
||||
show you how to get JAX running
|
||||
and provide you some tutorials to get started.
|
||||
|
||||
If looking to jump straight in take a look at the JAX quickstart.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
notebooks/quickstart
|
||||
|
||||
For most users starting out the key functionalities of JAX to become familiar with are :code:`jax.jit`,
|
||||
:code:`jax.grad`, and :code:`jax.vmap`.
|
||||
A good way to get familiar with this is with the Jax-101 tutorials.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
jax-101/index
|
||||
|
||||
If you're familiar with doing array-oriented computing with NumPy, you may find the following
|
||||
resources useful:
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
notebooks/thinking_in_jax
|
||||
notebooks/Common_Gotchas_in_JAX
|
||||
faq
|
||||
|
||||
If you prefer a video introduction here is one from JAX contributor Jake Vanderplas:
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<iframe width="640" height="360" src="https://www.youtube.com/embed/WdTeDXsOSj4"
|
||||
title="Intro to JAX: Accelerating Machine Learning research"
|
||||
frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture"
|
||||
allowfullscreen></iframe>
|
||||
@ -40,7 +40,7 @@ sys.path.insert(0, os.path.abspath('..'))
|
||||
# -- Project information -----------------------------------------------------
|
||||
|
||||
project = 'JAX'
|
||||
copyright = '2020, The JAX Authors. NumPy and SciPy documentation are copyright the respective authors.'
|
||||
copyright = '2023, The JAX Authors. NumPy and SciPy documentation are copyright the respective authors.'
|
||||
author = 'The JAX authors'
|
||||
|
||||
# The short X.Y version
|
||||
@ -72,6 +72,7 @@ extensions = [
|
||||
"sphinx_remove_toctrees",
|
||||
'sphinx_copybutton',
|
||||
'jax_extensions',
|
||||
'sphinx_design'
|
||||
]
|
||||
|
||||
intersphinx_mapping = {
|
||||
@ -174,6 +175,10 @@ html_favicon = '_static/favicon.png'
|
||||
# so a file named "default.css" will overwrite the builtin "default.css".
|
||||
html_static_path = ['_static']
|
||||
|
||||
html_css_files = [
|
||||
'style.css',
|
||||
]
|
||||
|
||||
# Custom sidebar templates, must be a dictionary that maps document names
|
||||
# to template names.
|
||||
#
|
||||
|
||||
17
docs/contributor_guide.rst
Normal file
17
docs/contributor_guide.rst
Normal file
@ -0,0 +1,17 @@
|
||||
.. _contributor-guide:
|
||||
|
||||
Developer Documentation
|
||||
=======================
|
||||
|
||||
JAX welcomes contributions from the community.
|
||||
See below for various install guides to get setup as a developer
|
||||
as well as developer-focused resources such as Jax Enhancement Proposals.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
contributing
|
||||
developer
|
||||
jax_internal_api
|
||||
autodidax
|
||||
jep/index
|
||||
151
docs/index.rst
151
docs/index.rst
@ -1,16 +1,76 @@
|
||||
JAX reference documentation
|
||||
===========================
|
||||
JAX: High-Performance Array Computing
|
||||
=====================================
|
||||
|
||||
JAX is Autograd_ and XLA_, brought together for high-performance numerical computing and machine learning research.
|
||||
It provides composable transformations of Python+NumPy programs: differentiate, vectorize,
|
||||
parallelize, Just-In-Time compile to GPU/TPU, and more.
|
||||
JAX is Autograd_ and XLA_, brought together for high-performance numerical computing.
|
||||
|
||||
.. grid:: 3
|
||||
|
||||
.. grid-item::
|
||||
|
||||
.. card:: Familiar API
|
||||
:class-card: key-ideas
|
||||
:shadow: None
|
||||
|
||||
JAX provides a familiar NumPy-style API for ease of adoption by researchers and engineers.
|
||||
|
||||
.. grid-item::
|
||||
|
||||
.. card:: Transformations
|
||||
:class-card: key-ideas
|
||||
:shadow: None
|
||||
|
||||
JAX includes composable function transformations for compilation, batching, automatic differentiation, and parallelization.
|
||||
|
||||
.. grid-item::
|
||||
|
||||
.. card:: Run Anywhere
|
||||
:class-card: key-ideas
|
||||
:shadow: None
|
||||
|
||||
The same code executes on multiple backends, including CPU, GPU, & TPU
|
||||
|
||||
.. note::
|
||||
JAX 0.4.1 introduces new parallelism APIs, including breaking changes to :func:`jax.experimental.pjit` and a new unified ``jax.Array`` type.
|
||||
Please see `Distributed arrays and automatic parallelization <https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html>`_ tutorial and the :ref:`jax-array-migration`
|
||||
guide for more information.
|
||||
|
||||
|
||||
.. grid:: 3
|
||||
|
||||
.. grid-item::
|
||||
|
||||
.. card:: :material-regular:`rocket_launch;2em` Getting Started
|
||||
:link: beginner-guide
|
||||
:link-type: ref
|
||||
:class-card: getting-started
|
||||
|
||||
.. grid-item::
|
||||
|
||||
.. card:: :material-regular:`library_books;2em` User Guides
|
||||
:link: user-guides
|
||||
:link-type: ref
|
||||
:class-card: user-guides
|
||||
|
||||
.. grid-item::
|
||||
|
||||
.. card:: :material-regular:`laptop_chromebook;2em` Developer Docs
|
||||
:link: contributor-guide
|
||||
:link-type: ref
|
||||
:class-card: developer-docs
|
||||
|
||||
Installation
|
||||
------------
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
pip install "jax[cpu]"
|
||||
|
||||
For installation on accelerators (GPU, TPU) and other installation options,
|
||||
see the `Install Guide`_ in the project README.
|
||||
|
||||
|
||||
.. toctree::
|
||||
:hidden:
|
||||
:maxdepth: 1
|
||||
:caption: Getting Started
|
||||
|
||||
@ -18,88 +78,27 @@ parallelize, Just-In-Time compile to GPU/TPU, and more.
|
||||
notebooks/quickstart
|
||||
notebooks/thinking_in_jax
|
||||
notebooks/Common_Gotchas_in_JAX
|
||||
faq
|
||||
|
||||
.. toctree::
|
||||
:hidden:
|
||||
:maxdepth: 1
|
||||
|
||||
jax-101/index
|
||||
|
||||
|
||||
.. toctree::
|
||||
:hidden:
|
||||
:maxdepth: 2
|
||||
:caption: Further Resources
|
||||
|
||||
debugging/index
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Reference Documentation
|
||||
|
||||
faq
|
||||
async_dispatch
|
||||
aot
|
||||
jaxpr
|
||||
notebooks/convolutions
|
||||
pytrees
|
||||
jax_array_migration
|
||||
type_promotion
|
||||
errors
|
||||
transfer_guard
|
||||
glossary
|
||||
changelog
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Advanced JAX Tutorials
|
||||
|
||||
notebooks/autodiff_cookbook
|
||||
multi_process
|
||||
notebooks/Distributed_arrays_and_automatic_parallelization
|
||||
notebooks/vmapped_log_probs
|
||||
notebooks/neural_network_with_tfds_data
|
||||
notebooks/Custom_derivative_rules_for_Python_code
|
||||
notebooks/How_JAX_primitives_work
|
||||
notebooks/Writing_custom_interpreters_in_Jax
|
||||
notebooks/Neural_Network_and_Data_Loading
|
||||
notebooks/xmap_tutorial
|
||||
notebooks/external_callbacks
|
||||
Custom_Operation_for_GPUs
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Developer documentation
|
||||
|
||||
contributing
|
||||
developer
|
||||
jax_internal_api
|
||||
autodidax
|
||||
jep/index
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: API documentation
|
||||
|
||||
user_guides
|
||||
advanced_guide
|
||||
contributor_guide
|
||||
notes
|
||||
jax
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Notes
|
||||
|
||||
api_compatibility
|
||||
deprecation
|
||||
concurrency
|
||||
gpu_memory_allocation
|
||||
profiling
|
||||
device_memory_profiling
|
||||
rank_promotion_warning
|
||||
|
||||
|
||||
Indices and tables
|
||||
==================
|
||||
|
||||
* :ref:`genindex`
|
||||
* :ref:`modindex`
|
||||
* :ref:`search`
|
||||
|
||||
|
||||
.. _Autograd: https://github.com/hips/autograd
|
||||
.. _XLA: https://www.tensorflow.org/xla
|
||||
.. _Install Guide: https://github.com/google/jax#installation
|
||||
@ -1,3 +1,7 @@
|
||||
:orphan:
|
||||
|
||||
.. _Jax-101:
|
||||
|
||||
Tutorial: JAX 101
|
||||
=================
|
||||
|
||||
|
||||
15
docs/notes.rst
Normal file
15
docs/notes.rst
Normal file
@ -0,0 +1,15 @@
|
||||
.. _notes:
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
api_compatibility
|
||||
deprecation
|
||||
concurrency
|
||||
gpu_memory_allocation
|
||||
profiling
|
||||
device_memory_profiling
|
||||
rank_promotion_warning
|
||||
@ -8,6 +8,7 @@ sphinx-book-theme>=0.3.3
|
||||
sphinx-copybutton>=0.5.0
|
||||
sphinx-remove-toctrees
|
||||
jupyter-sphinx>=0.3.2
|
||||
sphinx-design
|
||||
myst-nb
|
||||
|
||||
# Packages used for CI tests.
|
||||
|
||||
21
docs/user_guides.rst
Normal file
21
docs/user_guides.rst
Normal file
@ -0,0 +1,21 @@
|
||||
.. _user-guides:
|
||||
|
||||
User Guides
|
||||
===========
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
async_dispatch
|
||||
aot
|
||||
jaxpr
|
||||
notebooks/convolutions
|
||||
pytrees
|
||||
jax_array_migration
|
||||
type_promotion
|
||||
errors
|
||||
transfer_guard
|
||||
glossary
|
||||
changelog
|
||||
debugging/index
|
||||
Loading…
x
Reference in New Issue
Block a user