Ignore pytype errors produced with --use-functools-partial-overlay

PiperOrigin-RevId: 828850504

Author: superbobry

baselines/jft/vmoe.py

@@ -81,7 +81,7 @@ def load_checkpoint(config, mesh, efficient_ensemble_size):
   restore_checkpoint = functools.partial(
       partitioned.restore_checkpoint, tree=None, axis_resources=None)
   with mesh:
-    params = {p: restore_checkpoint(prefix=p) for p in model_init}
+    params = {p: restore_checkpoint(prefix=p) for p in model_init}  # pytype: disable=wrong-keyword-args
   return flax.core.freeze(params)
 
 

baselines/t5/data/tasks/toxic_comments.py

@@ -142,15 +142,15 @@ def _register_toxic_comments_classification_task(
     toxicity_label_threshold=TOXIC_LABEL_THRESHOLD)
 
 # Training data.
-_register_toxic_comments_task(
+_register_toxic_comments_task(  # pytype: disable=wrong-arg-types
     task_name='wikipedia_talk',
     tfds_name='wikipedia_toxicity_subtypes:0.3.1',
     tfds_splits=('train',),
 )
 
 # In-domain evaluation data. Note we will use test split for validation since
 # wikipedia_talk data does not have a validation split.
-_register_toxic_comments_task(
+_register_toxic_comments_task(  # pytype: disable=wrong-arg-types
     task_name='wikipedia_talk_eval_only',
     tfds_name='wikipedia_toxicity_subtypes:0.3.1',
     tfds_splits={
@@ -160,15 +160,15 @@ def _register_toxic_comments_classification_task(
 )
 
 # Out-of-domain evaluation data.
-_register_toxic_comments_task(
+_register_toxic_comments_task(  # pytype: disable=wrong-arg-types
     task_name='civil_comments_eval_only',
     tfds_name='civil_comments/CivilComments:1.1.2',
     tfds_splits=('validation', 'test'),
 )
 
 # Spurious correlation / tail generalization evaluation data, i.e., comments
 # with mentions of sensitive social / racial attributes.
-_register_toxic_comments_task(
+_register_toxic_comments_task(  # pytype: disable=wrong-arg-types
     task_name='civil_comments_identity_eval_only',
     tfds_name='civil_comments/CivilCommentsIdentities:1.1.2',
     tfds_splits=('validation', 'test'),
@@ -178,7 +178,7 @@ def _register_toxic_comments_classification_task(
 # contains different types of covert offensiveness (e.g., microaggression,
 # sarcasim, emoticons, etc).
 # Note this data only has positive label and does not have a validation split.
-_register_toxic_comments_task(
+_register_toxic_comments_task(  # pytype: disable=wrong-arg-types
     task_name='civil_comments_covert_eval_only',
     tfds_name='civil_comments/CivilCommentsCovert:1.1.2',
     tfds_splits={'validation': 'test'},