Merge pull request #3331 from levskaya:version_fix

PiperOrigin-RevId: 565109475

examples/lm1b/input_pipeline.py

@@ -322,7 +322,7 @@ def get_datasets(
     config: ml_collections.ConfigDict,
     *,
     n_devices: int,
-    vocab_path: Optional[str] = None
+    vocab_path: Optional[str] = None,
 ):
   """Load and return dataset of batched examples for use during training."""
   if vocab_path is None:

examples/ppo/ppo_lib.py

@@ -139,7 +139,7 @@ def train_step(
     *,
     clip_param: float,
     vf_coeff: float,
-    entropy_coeff: float
+    entropy_coeff: float,
 ):
   """Compilable train step.
 

examples/wmt/input_pipeline.py

@@ -47,7 +47,7 @@ def get_raw_dataset(
     dataset_builder: tfds.core.DatasetBuilder,
     split: str,
     *,
-    reverse_translation: bool = False
+    reverse_translation: bool = False,
 ) -> tf.data.Dataset:
   """Loads a raw WMT dataset and normalizes feature keys.
 
@@ -333,7 +333,7 @@ def get_wmt_datasets(
     *,
     n_devices: int,
     reverse_translation: bool = True,
-    vocab_path: Optional[str] = None
+    vocab_path: Optional[str] = None,
 ):
   """Load and return dataset of batched examples for use during training."""
   if vocab_path is None:

pyproject.toml

@@ -4,6 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flax"
+requires-python = ">=3.9"
 description = "Flax: A neural network library for JAX designed for flexibility"
 keywords = []
 authors = [
@@ -25,7 +26,9 @@ classifiers = [
     "Intended Audience :: Developers",
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: Apache Software License",
-    "Programming Language :: Python :: 3.7",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dynamic = ["version", "readme"]

tests/linen/linen_module_test.py

@@ -2442,16 +2442,18 @@ def _interceptor(f, args, kwargs, context):
 
     mod = IdentityModule(parent=None)
     x = 7
-    with nn.intercept_methods(
-        op_interceptor(lambda a: a + 1)
-    ), nn.intercept_methods(op_interceptor(lambda a: a**2)):
+    with (
+        nn.intercept_methods(op_interceptor(lambda a: a + 1)),
+        nn.intercept_methods(op_interceptor(lambda a: a**2)),
+    ):
       y = mod(x)
 
     self.assertEqual(y, (x**2) + 1)
 
-    with nn.intercept_methods(
-        op_interceptor(lambda a: a**2)
-    ), nn.intercept_methods(op_interceptor(lambda a: a + 1)):
+    with (
+        nn.intercept_methods(op_interceptor(lambda a: a**2)),
+        nn.intercept_methods(op_interceptor(lambda a: a + 1)),
+    ):
       y = mod(x)
 
     self.assertEqual(y, (x + 1) ** 2)