Replace # pytype: disable suppressions with # pyrefly: ignore

tensorflow_gnn/data/unigraph.py

@@ -45,7 +45,7 @@
 # Placeholder for optional Google-internal sorted string file format utils
 try:
   # pylint: disable-next=g-import-not-at-top
-  import google.cloud.bigquery_storage_v1 as bq_storage  # pytype: disable=import-error
+  import google.cloud.bigquery_storage_v1 as bq_storage  # pyrefly: ignore[missing-import]
 except ImportError:
   bq_storage = None
 
@@ -885,7 +885,7 @@ def iter_tfrecord_examples(
           Union[tfgnn.proto.NodeSet, tfgnn.proto.EdgeSet]] = None
       ) -> Iterable[Example]:
     """Yields `tf.Example` from tfrecord file."""
-    for example in tf.data.TFRecordDataset(file_path):  # pyrefly: ignore[bad-instantiation]
+    for example in tf.data.TFRecordDataset(file_path):
       yield Example.FromString(example.numpy())
 
   @staticmethod
@@ -899,7 +899,7 @@ def iter_csv_examples(
       converters = build_converter_from_schema(fset.features)
     else:
       converters = None
-    csv_records = csv.DictReader(gfile.GFile(file_path, "r"))  # pyrefly: ignore[bad-argument-type]
+    csv_records = csv.DictReader(gfile.GFile(file_path, "r"))
     for csv_record in csv_records:
       yield _csv_fields_to_example(csv_record.items(), converters=converters)
 
@@ -1066,7 +1066,7 @@ def iter_graph_via_schema(schema: tfgnn.GraphSchema) -> Dict[
           str, Union[
               Iterable[Tuple[bytes, bytes, Example]],
               Iterable[Tuple[bytes, Example]]]]]:
-    return {  # pyrefly: ignore[bad-return]
+    return {
         tfgnn.NODES: DictStreams.iter_nodes_via_schema(schema),  # pyrefly: ignore[bad-assignment]
         tfgnn.EDGES: DictStreams.iter_edges_via_schema(schema),  # pyrefly: ignore[bad-assignment]
     }

tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py

@@ -254,7 +254,7 @@ def test_sampling_stats(self):
     edges_layer = sampler.UniformEdgesSampler(
         sampler.KeyToTfExampleAccessor(
             sampler.InMemStringKeyToBytesAccessor(
-                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, bad-assignment]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-assignment]
             ),
             features_spec={
                 '#target': tf.TensorSpec([None], tf.string),
@@ -313,7 +313,7 @@ def test_missing_values(self):
     edges_layer = sampler.UniformEdgesSampler(
         sampler.KeyToTfExampleAccessor(
             sampler.InMemStringKeyToBytesAccessor(
-                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, bad-assignment]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-assignment]
             ),
             features_spec={
                 '#target': tf.TensorSpec([None], tf.int64),
@@ -437,7 +437,7 @@ def test_edge_features(self, feature_name: str, shape: List[int]):
     edges_layer = sampler.UniformEdgesSampler(
         sampler.KeyToTfExampleAccessor(
             sampler.InMemStringKeyToBytesAccessor(
-                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, bad-assignment]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-assignment]
             ),
             features_spec={
                 '#target': tf.TensorSpec([None], tf.string),
@@ -504,7 +504,7 @@ def test_ragged_features(self):
     edges_layer = sampler.UniformEdgesSampler(
         sampler.KeyToTfExampleAccessor(
             sampler.InMemStringKeyToBytesAccessor(
-                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, bad-assignment]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-assignment]
             ),
             features_spec={
                 'neighbors': tf.TensorSpec([None], tf.int64),

tensorflow_gnn/experimental/sampler/beam/executor_lib_test.py

@@ -214,16 +214,16 @@ def test_any_composite(self):
 
     i = tf.keras.Input([2], name='input')
     o = i
-    o = tf.keras.layers.Lambda(tf.sparse.from_dense)(o)  # pyrefly: ignore[not-callable]
-    o = tf.keras.layers.Lambda(tf.math.negative)(o)  # pyrefly: ignore[not-callable]
+    o = tf.keras.layers.Lambda(tf.sparse.from_dense)(o)
+    o = tf.keras.layers.Lambda(tf.math.negative)(o)
     # Add identity composite layer which splits eval data on five pieces:
     # <input> => <to negative sparse> => identity -> <to dense> => output.
-    o = Identity()(o)  # pyrefly: ignore[not-callable]
+    o = Identity()(o)
 
     def fn(t):
       return tf.sparse.to_dense(t)
 
-    o = tf.keras.layers.Lambda(fn)(o)  # pyrefly: ignore[not-callable]
+    o = tf.keras.layers.Lambda(fn)(o)
     model = tf.keras.Model(inputs=i, outputs=o)
     program, artifacts = sampler.create_program(model)
     self.assertLen(program.eval_dag.stages, 5)

tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py

@@ -57,7 +57,7 @@ def edge_sampler_factory(sampling_op):
       self.assertEqual(sampling_op.edge_set_name, 'a->b')
       if ids_dtype == tf.string:
         accessor = sampler.InMemStringKeyToBytesAccessor(
-            keys_to_values={b'a': b''}  # pyrefly: ignore[bad-argument-type, bad-assignment]
+            keys_to_values={b'a': b''}  # pyrefly: ignore[bad-assignment]
         )
       else:
         accessor = sampler.InMemIntegerKeyToBytesAccessor(

tensorflow_gnn/keras/keras_tensors.py

@@ -156,7 +156,7 @@ def handle(self, args, kwargs):
       )
 
     layer_args, layer_kwargs, args_index = _pack_args(args, kwargs)
-    return _GraphPieceClassMethod(  # pyrefly: ignore[not-callable]
+    return _GraphPieceClassMethod(
         cls._type_spec_cls(), self._cls_method_name, args_index
     )(layer_args, **layer_kwargs)
 
@@ -335,7 +335,7 @@ def __call__(self, *args, **kwargs):
       return self._op(*args, **kwargs)
 
     layer_args, layer_kwargs, args_index = _pack_args(args, kwargs)
-    return TFGNNOpLambda(self._resistry_key, args_index)(  # pyrefly: ignore[not-callable]
+    return TFGNNOpLambda(self._resistry_key, args_index)(
         layer_args, **layer_kwargs
     )
 

tensorflow_gnn/keras/layers/convolutions.py

@@ -145,6 +145,6 @@ def convolve(self, *,
     combined_input = ops.combine_values(inputs, self._combine_type)
 
     # Compute the result.
-    messages = self._message_fn(combined_input)  # pyrefly: ignore[not-callable]
+    messages = self._message_fn(combined_input)
     pooled_messages = pool_to_receiver(messages, reduce_type=self._reduce_type)
     return pooled_messages

tensorflow_gnn/keras/layers/item_dropout.py

@@ -74,4 +74,4 @@ def call(self, inputs):
     if self._dropout.noise_shape.rank != inputs.shape.rank:
       raise ValueError(f"Built for rank {self._dropout.noise_shape.rank}, "
                        f"called with input of rank {inputs.shape.rank}")
-    return self._dropout(inputs)  # pyrefly: ignore[not-callable]
+    return self._dropout(inputs)

tensorflow_gnn/keras/layers/next_state.py

@@ -141,7 +141,7 @@ def call(
       ]) -> const.FieldOrFields:
     net = tf.nest.flatten(inputs)
     net = tf.concat(net, axis=-1)
-    net = self._transformation(net)  # pyrefly: ignore[not-callable]
+    net = self._transformation(net)
     return net
 
 
@@ -239,7 +239,7 @@ def call(
     # Compute the state update.
     net = tf.nest.flatten(inputs)
     net = tf.concat(net, axis=-1)
-    net = self._residual_block(net)  # pyrefly: ignore[not-callable]
+    net = self._residual_block(net)
     if skip_connection_feature.shape[1:].num_elements() == 0:
       tf.get_logger().warning(
           "ResidualNextState() called on empty input state (latent node set?); "
@@ -253,7 +253,7 @@ def call(
           f"from {skip_connection_msg}.")
     else:
       net = tf.add(net, skip_connection_feature)
-    net = self._activation(net)  # pyrefly: ignore[not-callable]
+    net = self._activation(net)
     return net
 
 

tensorflow_gnn/models/gat_v2/layers.py

@@ -264,7 +264,7 @@ def convolve(self, *,
     # Form the attention query for each head.
     # [num_items, *extra_dims, num_heads, channels_per_head]
     assert receiver_input is not None, "__init__() should have checked this."
-    query = broadcast_from_receiver(self._split_heads(self._w_query(  # pyrefly: ignore[not-callable]
+    query = broadcast_from_receiver(self._split_heads(self._w_query(
         receiver_input)))
 
     # Form the attention value by transforming the configured inputs
@@ -286,7 +286,7 @@ def convolve(self, *,
 
     # Compute the attention logits and softmax to get the coefficients.
     # [num_items, *extra_dims, num_heads, 1]
-    logits = tf.expand_dims(self._attention_logits_fn(attention_features), -1)  # pyrefly: ignore[not-callable]
+    logits = tf.expand_dims(self._attention_logits_fn(attention_features), -1)
     attention_coefficients = extra_receiver_ops["softmax"](logits)  # pyrefly: ignore[unsupported-operation]
 
     if self._edge_dropout_layer is not None:
@@ -295,7 +295,7 @@ def convolve(self, *,
       # have the same effect as edge dropout.
       # Also, note that `keras.layers.Dropout` upscales the remaining values,
       # which should maintain the sum-up-to-1 per node in expectation.
-      attention_coefficients = self._edge_dropout_layer(attention_coefficients,  # pyrefly: ignore[not-callable]
+      attention_coefficients = self._edge_dropout_layer(attention_coefficients,
                                                         **kwargs)
 
     # Apply the attention coefficients to the transformed query.

tensorflow_gnn/models/gcn/gcn_conv.py

@@ -279,7 +279,7 @@ def get_degree(node_tag: tfgnn.IncidentNodeTag):
       else:
         pooled += normalized_values
 
-    return self._filter(pooled)  # pyrefly: ignore[not-callable]
+    return self._filter(pooled)
 
 
 def GCNHomGraphUpdate(*,  # To be called like a class initializer.  pylint: disable=invalid-name

tensorflow_gnn/models/graph_sage/layers.py

@@ -120,9 +120,9 @@ def convolve(self, *, sender_node_input: Optional[tf.Tensor],
     assert extra_receiver_ops is None, "Internal error: bad super().__init__()"
     assert sender_node_input is not None, "sender_node_input can't be None."
     result = broadcast_from_sender_node(sender_node_input)
-    result = self._dropout(result, training=training)  # pyrefly: ignore[not-callable]
+    result = self._dropout(result, training=training)
     result = pool_to_receiver(result, reduce_type=self._reduce_type)
-    result = self._transform_neighbor_fn(result)  # pyrefly: ignore[not-callable]
+    result = self._transform_neighbor_fn(result)
     return result
 
 
@@ -251,11 +251,11 @@ def convolve(self, *, sender_node_input: Optional[tf.Tensor],
     assert extra_receiver_ops is None, "Internal error: bad super().__init__()"
     assert sender_node_input is not None, "sender_node_input can't be None."
     result = broadcast_from_sender_node(sender_node_input)
-    result = self._dropout(result, training=training)  # pyrefly: ignore[not-callable]
+    result = self._dropout(result, training=training)
     # The "Pooling aggregator" from Eq. (3) of the paper, plus dropout.
-    result = self._pooling_transform_fn(result)  # pyrefly: ignore[not-callable]
+    result = self._pooling_transform_fn(result)
     result = pool_to_receiver(result, reduce_type=self._reduce_type)
-    result = self._transform_neighbor_fn(result)  # pyrefly: ignore[not-callable]
+    result = self._transform_neighbor_fn(result)
     return result
 
 
@@ -431,12 +431,12 @@ def call(self,
           tfgnn.reverse_tag(self._receiver_tag))
       sender_node_values = graph.node_sets[sender_node_set_name][
           self._sender_node_feature]
-      sender_node_values = self._dropout(sender_node_values)  # pyrefly: ignore[not-callable]
+      sender_node_values = self._dropout(sender_node_values)
       if not self._share_weights:
         sender_node_values = self._transform_edge_fn_dict[edge_set_name](
             sender_node_values)
       else:
-        sender_node_values = self._node_transform_fn(sender_node_values)  # pyrefly: ignore[not-callable]
+        sender_node_values = self._node_transform_fn(sender_node_values)
       broadcasted_sender_values = tfgnn.broadcast_node_to_edges(
           graph,
           edge_set_name,
@@ -463,8 +463,8 @@ def call(self,
     # aggregate with self node states only if add_self_loop is enabled.
     if self._add_self_loop:
       self_node_values = graph.node_sets[node_set_name][self._self_node_feature]
-      self_node_values = self._dropout(self_node_values)  # pyrefly: ignore[not-callable]
-      self_node_values = self._node_transform_fn(self_node_values)  # pyrefly: ignore[not-callable]
+      self_node_values = self._dropout(self_node_values)
+      self_node_values = self._node_transform_fn(self_node_values)
       pooled_node_states_list.append(self_node_values)
       if self._reduce_type == "mean":
         edge_set_in_degrees_list.append(tf.ones(total_size))
@@ -557,7 +557,7 @@ def convolutions_factory(edge_set_name, receiver_tag):
           sender_node_feature=feature_name,
           reduce_type=reduce_type,
           units=units,
-          hidden_units=hidden_units,  # pyrefly: ignore[bad-argument-type]
+          hidden_units=hidden_units,
           use_bias=use_bias,
           dropout_rate=dropout_rate,
           **kwargs)
@@ -739,8 +739,8 @@ def call(self, inputs, training):
     if unused_context_state:
       raise ValueError(
           "Input from context is not supported by GraphSAGENextState")
-    result = self._dropout(old_node_state, training=training)  # pyrefly: ignore[not-callable]
-    result = self._self_transform(result)  # pyrefly: ignore[not-callable]
+    result = self._dropout(old_node_state, training=training)
+    result = self._self_transform(result)
     result = [result, *[v for _, v in sorted(edge_inputs_dict.items())]]
     result = tfgnn.combine_values(result, self._combine_type)
     if self._use_bias:

tensorflow_gnn/models/mt_albis/layers.py

@@ -198,20 +198,20 @@ def call(
     # Collect and combine pooled messages (conv results) from edge sets.
     if edge_set_inputs:
       edge_input = self._combine_edge_inputs(edge_set_inputs)
-      edge_input = self._dropout(edge_input)  # pyrefly: ignore[not-callable]
+      edge_input = self._dropout(edge_input)
       flat_inputs.append(edge_input)
     # Collect a context input, if any. (Empty Mapping means none.)
     if isinstance(context_input, Mapping) and not context_input:
       pass
     else:
       context_input = _require_single_tensor(context_input,
                                              "input from context")
-      context_input = self._dropout(context_input)  # pyrefly: ignore[not-callable]
+      context_input = self._dropout(context_input)
       flat_inputs.append(context_input)
 
     net = tf.concat(flat_inputs, axis=-1)
-    net = self._dense(net)  # pyrefly: ignore[not-callable]
-    net = self._dropout(net)  # pyrefly: ignore[not-callable]
+    net = self._dense(net)
+    net = self._dropout(net)
 
     if (self._next_state_type == "residual" and
         input_state.shape[1:].num_elements() != 0):
@@ -223,7 +223,7 @@ def call(
             f"output shape {net.shape.as_list()}.")
       net = tf.add(net, input_state)
 
-    net = self._normalization(net)  # pyrefly: ignore[not-callable]
+    net = self._normalization(net)
     return net
 
   def _combine_edge_inputs(self, edge_set_inputs):

tensorflow_gnn/models/multi_head_attention/layers.py

@@ -377,11 +377,11 @@ def convolve(self,
                **kwargs) -> tf.Tensor:
 
     # Apply dropout on the inputs.
-    receiver_input = self._inputs_dropout_layer(receiver_input)  # pyrefly: ignore[not-callable]
+    receiver_input = self._inputs_dropout_layer(receiver_input)
     if sender_node_input is not None:
-      sender_node_input = self._inputs_dropout_layer(sender_node_input)  # pyrefly: ignore[not-callable]
+      sender_node_input = self._inputs_dropout_layer(sender_node_input)
     if sender_edge_input is not None:
-      sender_edge_input = self._inputs_dropout_layer(sender_edge_input)  # pyrefly: ignore[not-callable]
+      sender_edge_input = self._inputs_dropout_layer(sender_edge_input)
 
     # Determine the width of transformed queries and create transfomations.
     # If transform_keys is true, queries will be transformed to
@@ -417,7 +417,7 @@ def convolve(self,
     # [num_items, *extra_dims, num_heads, channels_per_head]
     # Otherwise, the shape is: [num_items, *extra_dims, num_heads, keys_width].
     assert receiver_input is not None, "__init__() should have checked this."
-    queries = self._w_query(receiver_input)  # pyrefly: ignore[not-callable]
+    queries = self._w_query(receiver_input)
     queries = self._attention_activation(queries)
     queries = broadcast_from_receiver(self._split_query_heads(queries))
 
@@ -491,7 +491,7 @@ def convolve(self,
     # should have the same effect as edge dropout. Also, note that
     # `keras.layers.Dropout` upscales the remaining values, which should
     # maintain the sum-up-to-1 per node in expectation.
-    attention_coefficients = self._edge_dropout_layer(attention_coefficients,  # pyrefly: ignore[not-callable]
+    attention_coefficients = self._edge_dropout_layer(attention_coefficients,
                                                       **kwargs)
 
     # Compute the pooled values by
@@ -541,7 +541,7 @@ def convolve(self,
       pooled_inputs = pool_to_receiver(weighted_inputs, reduce_type="sum")
       # Apply the transformation.
       # [num_receivers, *extra_dims, num_heads, per_head_channels]
-      pooled_values = self._w_sender_pooled_to_value(pooled_inputs)  # pyrefly: ignore[not-callable]
+      pooled_values = self._w_sender_pooled_to_value(pooled_inputs)
 
     # Apply the nonlinearity on the final result.
     pooled_values = self._activation(pooled_values)