Improve documentation:

(1) add more structure, e.g. list clickable layers with detailed signature on separate pages, instead of a monolithic wall of text. (2) Add drop-down TOC on the left, and some simple sections like linear / nonlinear layers etc. (3) Type rendering is improved following prior change introducting typing protocols. (4) All sphinx warnings are fixed, so we can now `fail_on_warning`, which will prevent silent errors / empty docs happening in the past. PiperOrigin-RevId: 429647336
google · Feb 18, 2022 · 22b3aae · 22b3aae
1 parent b0d489e
commit 22b3aae
Show file tree

Hide file tree

Showing 14 changed files with 218 additions and 39 deletions.
diff --git a/.readthedocs.yml b/.readthedocs.yml
@@ -8,6 +8,7 @@ version: 2
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
   configuration: docs/conf.py
+  fail_on_warning: True
 
 # Optionally build your docs in additional formats such as PDF and ePub
 formats:

diff --git a/docs/batching.rst b/docs/batching.rst
@@ -0,0 +1,8 @@
+Batching -- using multiple devices
+============================================================
+
+.. default-role:: code
+.. automodule:: neural_tangents._src.batching
+.. automodule:: neural_tangents
+  :noindex:
+  :members: batch
diff --git a/docs/conf.py b/docs/conf.py
@@ -46,6 +46,8 @@
     'sphinx_autodoc_typehints',
 ]
 
+# set_type_checking_flag = False
+
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['_templates']
 
@@ -97,7 +99,7 @@
 # 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']
+# html_static_path = ['_static']
 
 # Custom sidebar templates, must be a dictionary that maps document names
 # to template names.
@@ -170,6 +172,17 @@
      'Miscellaneous'),
 ]
 
+# add_module_names = False
+
+
+def remove_module_docstring(app, what, name, obj, options, lines):
+  if what == "module" and name == "neural_tangents":
+    del lines[:]
+
+
+def setup(app):
+  app.connect("autodoc-process-docstring", remove_module_docstring)
+
 
 # -- Options for Epub output -------------------------------------------------
 

diff --git a/docs/empirical.rst b/docs/empirical.rst
@@ -0,0 +1,29 @@
+.. default-role:: code
+
+Empirical -- finite NNGP and NTK
+======================================
+
+.. automodule:: neural_tangents._src.empirical
+.. currentmodule:: neural_tangents
+
+Kernel functions
+--------------------------------------
+Finite-width NNGP and/or NTK kernel functions.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    empirical_kernel_fn
+    empirical_nngp_fn
+    empirical_ntk_fn
+
+
+Linearization and Taylor expansion
+--------------------------------------
+Decorators to Taylor-expand around function parameters.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    linearize
+    taylor_expand
diff --git a/docs/index.rst b/docs/index.rst
@@ -8,11 +8,11 @@ neural networks.
    :maxdepth: 2
    :caption: Topics:
 
-   neural_tangents.stax
-   neural_tangents._src.empirical
-   neural_tangents.predict
-   neural_tangents._src.batching
-   neural_tangents._src.monte_carlo
+   stax
+   empirical
+   predict
+   batching
+   monte_carlo
 
 Indices and tables
 ==================

diff --git a/docs/monte_carlo.rst b/docs/monte_carlo.rst
@@ -0,0 +1,7 @@
+Monte Carlo Sampling
+======================================
+
+.. default-role:: code
+.. automodule:: neural_tangents._src.monte_carlo
+.. automodule:: neural_tangents
+  :members: monte_carlo_kernel_fn
diff --git a/docs/neural_tangents.batching.rst b/docs/neural_tangents.batching.rst
diff --git a/docs/neural_tangents.empirical.rst b/docs/neural_tangents.empirical.rst
diff --git a/docs/neural_tangents.monte_carlo.rst b/docs/neural_tangents.monte_carlo.rst
diff --git a/docs/neural_tangents.predict.rst b/docs/neural_tangents.predict.rst
diff --git a/docs/neural_tangents.stax.rst b/docs/neural_tangents.stax.rst
diff --git a/docs/predict.rst b/docs/predict.rst
@@ -0,0 +1,39 @@
+.. default-role:: code
+
+Predict -- inference with NNGP and NTK or linearized networks
+=============================================================
+
+.. automodule:: neural_tangents._src.predict
+.. currentmodule:: neural_tangents.predict
+
+
+Prediction / inference functions
+--------------------------------------
+Functions to make train/test set predictions given NNGP/NTK kernels or the linearized function.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    gradient_descent
+    gradient_descent_mse
+    gradient_descent_mse_ensemble
+    gp_inference
+
+
+Utilities
+--------------------------------------
+.. autosummary::
+ :toctree: _autosummary
+
+    max_learning_rate
+
+
+Helper classes
+--------------------------------------
+Dataclasses and namedtuples used to return predictions.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    Gaussian
+    ODEState
diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -5,6 +5,3 @@ jaxlib
 ipykernel
 nbsphinx
 sphinx-autodoc-typehints
-# The next packages are for notebooks
-matplotlib
-.
diff --git a/docs/stax.rst b/docs/stax.rst
@@ -0,0 +1,115 @@
+.. default-role:: code
+
+Stax -- infinite networks (NNGP, NTK)
+======================================
+
+.. automodule:: neural_tangents.stax
+
+
+Combinators
+--------------------------------------
+Layers to combine multiple other layers into one.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    serial
+    parallel
+
+
+Branching
+--------------------------------------
+Layers to split outputs into many, or combine many into ones.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    FanOut
+    FanInConcat
+    FanInProd
+    FanInSum
+
+
+Linear parametric
+--------------------------------------
+Linear layers with trainable parameters.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    Dense
+    Conv
+    ConvLocal
+    ConvTranspose
+    GlobalSelfAttention
+
+
+Linear nonparametric
+--------------------------------------
+Linear layers without any trainable parameters.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    Aggregate
+    AvgPool
+    Identity
+    DotGeneral
+    Dropout
+    Flatten
+    GlobalAvgPool
+    GlobalSumPool
+    ImageResize
+    LayerNorm
+    SumPool
+
+
+Elementwise nonlinear
+--------------------------------------
+Pointwise nonlinear layers.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    ABRelu
+    Abs
+    Cos
+    Elementwise
+    ElementwiseNumerical
+    Erf
+    Exp
+    ExpNormalized
+    Gaussian
+    Gelu
+    Hermite
+    LeakyRelu
+    Rbf
+    Relu
+    Sigmoid_like
+    Sign
+    Sin
+
+
+Helper enums
+--------------------------------------
+Enums for specifying layer properties. Strings can be used in their place.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    Padding
+    PositionalEmbedding
+
+
+For developers
+--------------------------------------
+Classes and decorators helpful for constructing your own layers.
+
+.. autosummary::
+ :toctree: _autosummary
+
+    layer
+    supports_masking
+    requires
+    Bool
+    Diagonal