Internal change

PiperOrigin-RevId: 435732468

official/nlp/data/dual_encoder_dataloader.py

@@ -124,7 +124,7 @@ class DualEncoderDataLoader(data_loader.DataLoader):
       raise ValueError('Expected {} to start with {}'.format(string, old))
 
     def _switch_key_prefix(d, old, new):
-      return {_switch_prefix(key, old, new): value for key, value in d.items()}
+      return {_switch_prefix(key, old, new): value for key, value in d.items()}  # pytype: disable=attribute-error  # trace-all-classes
 
     model_inputs = _switch_key_prefix(
         self._bert_tokenize(record, self._left_text_fields),

official/vision/beta/modeling/backbones/revnet.py

@@ -208,7 +208,7 @@ class RevNet(tf.keras.Model):
   @property
   def output_specs(self) -> Dict[int, tf.TensorShape]:
     """A dict of {level: TensorShape} pairs for the model output."""
-    return self._output_specs
+    return self._output_specs  # pytype: disable=bad-return-type  # trace-all-classes
 
 
 @factory.register_backbone_builder('revnet')

official/vision/beta/modeling/decoders/aspp.py

@@ -103,7 +103,7 @@ class ASPP(tf.keras.layers.Layer):
     if self._config_dict['pool_kernel_size']:
       pool_kernel_size = [
           int(p_size // 2**self._config_dict['level'])
-          for p_size in self._config_dict['pool_kernel_size']
+          for p_size in self._config_dict['pool_kernel_size']  # pytype: disable=attribute-error  # trace-all-classes
       ]
 
     self.aspp = self._aspp_layer(

official/vision/modeling/backbones/revnet.py

@@ -208,7 +208,7 @@ class RevNet(tf.keras.Model):
   @property
   def output_specs(self) -> Dict[int, tf.TensorShape]:
     """A dict of {level: TensorShape} pairs for the model output."""
-    return self._output_specs
+    return self._output_specs  # pytype: disable=bad-return-type  # trace-all-classes
 
 
 @factory.register_backbone_builder('revnet')

official/vision/modeling/decoders/aspp.py

@@ -103,7 +103,7 @@ class ASPP(tf.keras.layers.Layer):
     if self._config_dict['pool_kernel_size']:
       pool_kernel_size = [
           int(p_size // 2**self._config_dict['level'])
-          for p_size in self._config_dict['pool_kernel_size']
+          for p_size in self._config_dict['pool_kernel_size']  # pytype: disable=attribute-error  # trace-all-classes
       ]
 
     self.aspp = self._aspp_layer(