Update readthedocs requirements to include all neural tangents build …

…dependencies. Fix pytype python version PiperOrigin-RevId: 557183101
google · Aug 24, 2023 · b3b728e · b3b728e
1 parent db53e46
commit b3b728e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/.github/workflows/linux.yml b/.github/workflows/linux.yml
@@ -24,7 +24,7 @@ jobs:
         python-version: [3.9, '3.10', 3.11]
         JAX_ENABLE_X64: [0, 1]
 
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
 
     steps:
 

diff --git a/tests/predict_test.py b/tests/predict_test.py
@@ -300,7 +300,7 @@ def testNTKMeanCovPrediction(
     self.assertGreater(np.min(np.linalg.eigh(cov_train_inf)[0]), -1e-8)
 
     _kernel_fn = nt.empirical_kernel_fn(f)
-    kernel_fn = lambda x1, x2, params: _kernel_fn(x1, x2, 'ntk', params)
+    kernel_fn = jit(lambda x1, x2, params: _kernel_fn(x1, x2, 'ntk', params))
 
     def predict_empirical(key):
       _, params = init_fn(key, train_shape)

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -19,7 +19,7 @@
 import logging
 import os
 from types import ModuleType
-from typing import Callable, Optional, Sequence
+from typing import Callable,  Optional, Sequence
 
 from absl import flags
 from absl.testing import parameterized
@@ -41,7 +41,7 @@
 
 flags.DEFINE_integer(
     'nt_num_generated_cases',
-    int(os.getenv('NT_NUM_GENERATED_CASES', '2')),
+    int(os.getenv('NT_NUM_GENERATED_CASES', '4')),
     help='Number of generated cases to test'
 )