Add note to docs describing how pytree arguments work. (google#3284)

Addresses google#3095. I'm not sure if we wanna link to this from API docstrings. This also subsumes the original pytrees notebook.
majnemer · Jun 3, 2020 · 66ba734 · 66ba734
1 parent f8bab4a
commit 66ba734
Show file tree

Hide file tree

Showing 3 changed files with 252 additions and 330 deletions.
diff --git a/docs/index.rst b/docs/index.rst
@@ -23,7 +23,6 @@ For an introduction to JAX, start at the
 
    notebooks/Common_Gotchas_in_JAX
    notebooks/Custom_derivative_rules_for_Python_code
-   notebooks/JAX_pytrees
    notebooks/XLA_in_Python
    notebooks/How_JAX_primitives_work
    notebooks/Writing_custom_interpreters_in_Jax.ipynb
@@ -39,6 +38,7 @@ For an introduction to JAX, start at the
    concurrency
    gpu_memory_allocation
    profiling
+   pytrees
    rank_promotion_warning
    type_promotion
 

diff --git a/docs/notebooks/JAX_pytrees.ipynb b/docs/notebooks/JAX_pytrees.ipynb