Add missing TF-related mock imports to autodoc.

PiperOrigin-RevId: 456339616
google · Jun 24, 2022 · 8f64046 · 8f64046
1 parent 1e4754f
commit 8f64046
Show file tree

Hide file tree

Showing 12 changed files with 2,986 additions and 3,559 deletions.
diff --git a/README.md b/README.md
diff --git a/docs/empirical.rst b/docs/empirical.rst
@@ -19,15 +19,6 @@ Finite-width NNGP and/or NTK kernel functions.
     empirical_nngp_fn
     empirical_ntk_fn
 
-NTK implementation
---------------------------------------
-An `IntEnum` specifying NTK implementation method.
-
-.. autosummary::
- :toctree: _autosummary
-
-    NtkImplementation
-
 
 Linearization and Taylor expansion
 --------------------------------------

diff --git a/notebooks/Disentangling_Trainability_and_Generalization.ipynb b/notebooks/Disentangling_Trainability_and_Generalization.ipynb
@@ -42,9 +42,9 @@
    },
    "outputs": [],
    "source": [
-        "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
-        "!pip install -q git+https://www.github.com/google/neural-tangents"
+    "!pip install --upgrade pip\n",
+    "!pip install --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
+    "!pip install -q git+https://www.github.com/google/neural-tangents"
    ]
   },
   {