Fix a pytype error, add Index to readthedocs, reorder doc entries a…

…lphabetically, minor readme changes.

PiperOrigin-RevId: 460585349

README.md

@@ -1,10 +1,10 @@
 # **Stand with Ukraine!** 🇺🇦
 
-Freedom of thought is fundamental to all of science. Right now, our freedom is being suppressed with carpet bombing of civilians in Ukraine. **Don't be against the war - fight against the war!** Support Ukraine at **[stopputin.net](https://stopputin.net/)**.
+Freedom of thought is fundamental to all of science. Right now, our freedom is being suppressed with carpet bombing of civilians in Ukraine. **Don't be against the war - fight against the war! [supportukrainenow.org](https://supportukrainenow.org/)**.
 
 ### News
 
-Our paper "[Fast Finite Width Neural Tangent Kernel](https://arxiv.org/abs/2206.08720)" has been accepted to ICML2022, and the respective code [submitted](https://github.com/romanngg/neural-tangents/commit/60b6c16758652f4526536409b5bc90602287b868) (available starting from version `0.6.0`).
+**[Fast Finite Width Neural Tangent Kernel](https://arxiv.org/abs/2206.08720)** is accepted to ICML2022, and the respective code [submitted](https://github.com/romanngg/neural-tangents/commit/60b6c16758652f4526536409b5bc90602287b868) (available from `0.6.0`). See [video](https://youtu.be/8MWOhYg89fY?t=10984) and [example](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/empirical_ntk_fcn.ipynb).
 
 # Neural Tangents
 [**ICLR 2020 Video**](https://iclr.cc/virtual_2020/poster_SklD9yrFPS.html)
@@ -61,7 +61,7 @@ An easy way to get started with Neural Tangents is by playing around with the fo
 - [Function Space Linearization](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/function_space_linearization.ipynb)
 - [Neural Network Phase Diagram](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/phase_diagram.ipynb)
 - [Performance Benchmark](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/myrtle_kernel_with_neural_tangents.ipynb): simple benchmark for Myrtle kernels used in [[16]](#16-neural-kernels-without-tangents). Also see [Performance](#myrtle-network)
-- Empirical NTK:
+- [**New**] Empirical NTK:
   - [Fully-connected network](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/empirical_ntk_fcn.ipynb)
   - [FLAX ResNet18](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/empirical_ntk_resnet.ipynb)
   - [Experimental: Tensorflow ResNet50](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/experimental/empirical_ntk_resnet_tf.ipynb)

docs/index.rst

@@ -22,9 +22,9 @@ neural networks (a.k.a. NTK, NNGP).
    :maxdepth: 2
    :caption: Internal:
 
+   experimental
    kernel
    typing
-   experimental
 
 .. toctree::
    :maxdepth: 2

docs/predict.rst

@@ -16,10 +16,10 @@ Functions to make train/test set predictions given NNGP/NTK kernels or the linea
 .. autosummary::
  :toctree: _autosummary
 
+    gp_inference
     gradient_descent
     gradient_descent_mse
     gradient_descent_mse_ensemble
-    gp_inference
 
 
 Utilities

docs/stax.rst

@@ -16,8 +16,8 @@ Layers to combine multiple other layers into one.
 .. autosummary::
  :toctree: _autosummary
 
-    serial
     parallel
+    serial
 
 
 Branching
@@ -27,10 +27,10 @@ Layers to split outputs into many, or combine many into ones.
 .. autosummary::
  :toctree: _autosummary
 
-    FanOut
     FanInConcat
     FanInProd
     FanInSum
+    FanOut
 
 
 Linear parametric
@@ -40,10 +40,10 @@ Linear layers with trainable parameters.
 .. autosummary::
  :toctree: _autosummary
 
-    Dense
     Conv
     ConvLocal
     ConvTranspose
+    Dense
     GlobalSelfAttention
 
 
@@ -56,13 +56,14 @@ Linear layers without any trainable parameters.
 
     Aggregate
     AvgPool
-    Identity
     DotGeneral
     Dropout
     Flatten
     GlobalAvgPool
     GlobalSumPool
+    Identity
     ImageResize
+    Index
     LayerNorm
     SumPool
 
@@ -114,8 +115,8 @@ Classes and decorators helpful for constructing your own layers.
 .. autosummary::
  :toctree: _autosummary
 
-    layer
-    supports_masking
-    requires
     Bool
     Diagonal
+    layer
+    requires
+    supports_masking

neural_tangents/_src/stax/linear.py

@@ -2794,12 +2794,6 @@ def Index(
 ) -> InternalLayerMasked:
   """Index into the array mimicking :cls:`onp.ndarray` indexing.
 
-  .. warning::
-    Two limitations in the kernel regime (`kernel_fn`): the `channel_axis`
-    (infinite width) cannot be indexed, and the `batch_axis` can only be
-    indexed with tuples/slices, but not integers, since the library requires
-    there always to be a batch axis in a `Kernel`.
-
   Args:
     idx:
       a `slice` object that would result from indexing an array as `x[idx]`.
@@ -2817,6 +2811,22 @@ def Index(
   Returns:
     `(init_fn, apply_fn, kernel_fn)`.
 
+  Raises:
+    NotImplementedError:
+      If the `channel_axis` (infinite width) is indexed
+      (except for `:` or `...`) in the kernel regime (`kernel_fn`).
+
+    NotImplementedError:
+      If the `batch_axis` is indexed with an integer (as opposed to a tuple or
+      slice) in the kernel regime (`kernel_fn`), since the library currently
+      requires there always to be `batch_axis` in  the kernel regime (while
+      indexing with integers removes the respective axis).
+
+    ValueError:
+      If `init_fn` is called on a shape with dummy axes (with sizes like `-1`
+      or `None`), that are indexed with non-trivial (not `:` or `...`) slices.
+      For indexing, the size of the respective axis needs to be specified.
+
   Example:
     >>> from neural_tangents import stax
     >>> #

neural_tangents/_src/utils/kernel.py

@@ -477,8 +477,8 @@ def __getitem__(self, idx: utils.SliceType) -> 'Kernel':
     cov_slice = cov_batch_slice + spatial_idx
 
     nngp = self.nngp[nngp_slice]
-    ntk = (self.ntk if (self.ntk is None or self.ntk.ndim == 0) else
-           self.ntk[nngp_slice])  # pytype: disable=attribute-error
+    ntk = (self.ntk if (self.ntk is None or self.ntk.ndim == 0) else  # pytype: disable=attribute-error
+           self.ntk[nngp_slice])
 
     cov1 = self.cov1[cov_slice]
     cov2 = None if self.cov2 is None else self.cov2[cov_slice]