Support python 3.12 and misc

.github/workflows/linux.yml

@@ -21,7 +21,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.9, '3.10', 3.11]
+        python-version: ['3.10', 3.11, 3.12]
         JAX_ENABLE_X64: [0, 1]
 
     runs-on: ubuntu-latest
@@ -32,10 +32,10 @@ jobs:
       uses: styfle/[email protected]
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v4.1.1
+    - uses: actions/checkout@v4.2.1
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5.0.0
+      uses: actions/setup-python@v5.2.0
       with:
         python-version: ${{ matrix.python-version }}
 
@@ -54,7 +54,7 @@ jobs:
         JAX_ENABLE_X64=${{ matrix.JAX_ENABLE_X64 }} PYTHONHASHSEED=0 pytest -n auto --cov=neural_tangents --cov-report=xml --cov-report=term
 
     - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@v4.0.1
+      uses: codecov/codecov-action@v4.6.0
       with:
         file: ./coverage.xml
 

.github/workflows/macos.yml

@@ -21,7 +21,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.9, '3.10', 3.11]
+        python-version: ['3.10', 3.11, 3.12]
         JAX_ENABLE_X64: [0]
 
     runs-on: macos-latest
@@ -32,10 +32,10 @@ jobs:
       uses: styfle/[email protected]
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v4.1.1
+    - uses: actions/checkout@v4.2.1
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5.0.0
+      uses: actions/setup-python@v5.2.0
       with:
         python-version: ${{ matrix.python-version }}
 
@@ -54,7 +54,7 @@ jobs:
         JAX_ENABLE_X64=${{ matrix.JAX_ENABLE_X64 }} PYTHONHASHSEED=0 pytest -n auto --cov=neural_tangents --cov-report=xml --cov-report=term
 
     - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@v4.0.1
+      uses: codecov/codecov-action@v4.6.0
       with:
         file: ./coverage.xml
 

.github/workflows/pytype.yml

@@ -27,10 +27,10 @@ jobs:
       uses: styfle/[email protected]
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v4.1.1
+    - uses: actions/checkout@v4.2.1
 
     - name: Set up Python 3.10
-      uses: actions/setup-python@v5.0.0
+      uses: actions/setup-python@v5.2.0
       with:
         python-version: '3.10'
 

.github/workflows/sketching.yml

@@ -31,10 +31,10 @@ jobs:
       uses: styfle/[email protected]
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v4.1.1
+    - uses: actions/checkout@v4.2.1
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5.0.0
+      uses: actions/setup-python@v5.2.0
       with:
         python-version: ${{ matrix.python-version }}
 
@@ -53,7 +53,7 @@ jobs:
         JAX_ENABLE_X64=${{ matrix.JAX_ENABLE_X64 }} PYTHONHASHSEED=0 pytest experimental/tests/ -n auto --cov=experimental/ --cov-report=xml --cov-report=term
 
     - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@v4.0.1
+      uses: codecov/codecov-action@v4.6.0
       with:
         file: ./coverage.xml
 

README.md

@@ -104,7 +104,7 @@ from jax.example_libraries import stax
 init_fn, apply_fn = stax.serial(
     stax.Dense(512), stax.Relu,
     stax.Dense(512), stax.Relu,
-    stax.Dense(1)
+    stax.Dense(1),
 )
 
 key = random.PRNGKey(1)
@@ -123,7 +123,7 @@ from neural_tangents import stax
 init_fn, apply_fn, kernel_fn = stax.serial(
     stax.Dense(512), stax.Relu(),
     stax.Dense(512), stax.Relu(),
-    stax.Dense(1)
+    stax.Dense(1),
 )
 
 key1, key2 = random.split(random.PRNGKey(1))

examples/datasets.py

@@ -44,7 +44,7 @@ def get_dataset(
     permute_train=False,
     do_flatten_and_normalize=True,
     data_dir=None,
-    input_key='image'
+    input_key='image',
 ):
   """Download, parse and process a dataset to unit scale and one-hot labels."""
   # Need this following http://cl/378185881 to prevent GPU test breakages.
@@ -133,9 +133,10 @@ def embed_glove(xs, glove_path, max_sentence_length=1000, mask_constant=1000.):
   xs = list(map(_decode, xs))
   tokenizer = tf.keras.preprocessing.text.Tokenizer()
   tokenizer.fit_on_texts(np.concatenate(xs))
-  glove_embedding_layer = _get_glove_embedding_layer(tokenizer,
-                                                     glove_path,
-                                                     max_sentence_length)
+  glove_embedding_layer = _get_glove_embedding_layer(
+      tokenizer,
+      glove_path,
+  )
 
   def embed(x):
     # Replace strings with sequences of integer tokens.
@@ -147,7 +148,8 @@ def embed(x):
         x_tok,
         max_sentence_length,
         padding='post',
-        truncating='post')
+        truncating='post',
+    )
 
     # Replace integer tokens with word embeddings.
     x_emb = glove_embedding_layer(x_tok).numpy()
@@ -160,7 +162,7 @@ def embed(x):
   return map(embed, xs)
 
 
-def _get_glove_embedding_layer(tokenizer, glove_path, max_sentence_length):
+def _get_glove_embedding_layer(tokenizer, glove_path):
   """Get a Keras embedding layer for a given GloVe embeddings.
 
   Adapted from https://keras.io/examples/pretrained_word_embeddings/.
@@ -172,9 +174,6 @@ def _get_glove_embedding_layer(tokenizer, glove_path, max_sentence_length):
     glove_path:
       path to the GloVe embedding file.
 
-    max_sentence_length:
-      pad/truncate embeddings to this length.
-
   Returns:
     Keras embedding layer for a given GloVe embeddings.
   """
@@ -212,8 +211,8 @@ def _get_glove_embedding_layer(tokenizer, glove_path, max_sentence_length):
     embedding_layer = tf.keras.layers.Embedding(
         num_words, embedding_dim,
         embeddings_initializer=tf.keras.initializers.Constant(emb_mat),
-        input_length=max_sentence_length,
-        trainable=False)
+        trainable=False,
+    )
 
   return embedding_layer
 

examples/function_space.py

@@ -57,7 +57,7 @@ def main(unused_argv):
   opt_init, opt_apply, get_params = optimizers.sgd(_LEARNING_RATE)
   state = opt_init(params)
 
-  # Create an mse loss function and a gradient function.
+  # Create a mse loss function and a gradient function.
   loss = lambda fx, y_hat: 0.5 * jnp.mean((fx - y_hat) ** 2)
   grad_loss = jit(grad(lambda params, x, y: loss(apply_fn(params, x), y)))
 

examples/imdb.py

@@ -99,7 +99,7 @@ def main(*args, use_dummy_data: bool = False, **kwargs) -> None:
 
 
 def _get_dummy_data(
-    mask_constant: float
+    mask_constant: float,
 ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
   """Return dummy data for when downloading embeddings is not feasible."""
   n_train, n_test = 6, 6