No public description

PiperOrigin-RevId: 941349341

Author: h-joo

tensorflow_gnn/experimental/sampler/beam/accessors_test.py

@@ -68,7 +68,7 @@ def test_string_keys(self):
     }
     layer = sampler.KeyToTfExampleAccessor(
         sampler.InMemStringKeyToBytesAccessor(
-            keys_to_values=table, name='table'
+            keys_to_values=table, name='table'  # pyrefly: ignore[bad-argument-type]
         ),
         features_spec={
             's': tf.TensorSpec([], tf.int64),

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'
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type]
             ),
             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'
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type]
             ),
             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'
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type]
             ),
             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'
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type]
             ),
             features_spec={
                 'neighbors': tf.TensorSpec([None], tf.int64),

tensorflow_gnn/experimental/sampler/beam/executor_lib.py

@@ -268,11 +268,11 @@ def process(
         self, inputs: Tuple[bytes, Tuple[str, Values]]
     ) -> Iterator[Tuple[bytes, Values]]:
       example_id, (stage_id, values) = inputs
-      inputs = []
+      inputs = []  # pyrefly: ignore[bad-assignment]
       for matcher in self._stage.input_matchers:
         assert matcher.stage_id == stage_id, example_id
-        inputs.append(values[matcher.output_index])
-      yield (example_id, inputs)
+        inputs.append(values[matcher.output_index])  # pyrefly: ignore[missing-attribute]
+      yield (example_id, inputs)  # pyrefly: ignore[invalid-yield]
 
   @beam_typehints.with_input_types(
       Tuple[ExampleId, Iterable[Tuple[str, Values]]]
@@ -289,11 +289,11 @@ def process(
     ) -> Iterator[Tuple[bytes, Values]]:
       example_id, outputs = inputs
       outputs = dict(outputs)
-      inputs = []
+      inputs = []  # pyrefly: ignore[bad-assignment]
       for matcher in self._stage.input_matchers:
         values = outputs[matcher.stage_id]
-        inputs.append(values[matcher.output_index])
-      yield (example_id, inputs)
+        inputs.append(values[matcher.output_index])  # pyrefly: ignore[missing-attribute]
+      yield (example_id, inputs)  # pyrefly: ignore[invalid-yield]
 
   def __init__(self, stage: pb.Stage):
     self._stage = stage
@@ -364,11 +364,11 @@ def __init__(
   def expand(
       self, inputs: Tuple[Dict[str, PValues], Dict[str, PFeed]]
   ) -> PValues:
-    inputs, feeds = inputs
+    inputs, feeds = inputs  # pyrefly: ignore[bad-assignment]
     return _execute(
         self._eval_dag,
         self._layers,
-        inputs,
+        inputs,  # pyrefly: ignore[bad-argument-type]
         feeds=feeds,
         artifacts_path=self._artifacts_path,
     )

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)
-    o = tf.keras.layers.Lambda(tf.math.negative)(o)
+    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]
     # Add identity composite layer which splits eval data on five pieces:
     # <input> => <to negative sparse> => identity -> <to dense> => output.
-    o = Identity()(o)
+    o = Identity()(o)  # pyrefly: ignore[not-callable]
 
     def fn(t):
       return tf.sparse.to_dense(t)
 
-    o = tf.keras.layers.Lambda(fn)(o)
+    o = tf.keras.layers.Lambda(fn)(o)  # pyrefly: ignore[not-callable]
     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/sampler.py

@@ -135,7 +135,7 @@ def node_features_accessor_factory(
       node_set_name: tfgnn.NodeSetName,
   ) -> sampler.KeyToTfExampleAccessor:
     if not graph_schema.node_sets[node_set_name].features:
-      return None
+      return None  # pyrefly: ignore[bad-return]
 
     node_features = graph_schema.node_sets[node_set_name].features
 

tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py

@@ -44,7 +44,7 @@ def _get_sampling_model(self, ids_dtype) -> tf.keras.Model:
     graph_schema.node_sets['a'].description = 'node set b'
     graph_schema.edge_sets['a->b'].source = 'a'
     graph_schema.edge_sets['a->b'].target = 'b'
-    graph_schema.edge_sets['a->a'].features['weights'].dtype = 1
+    graph_schema.edge_sets['a->a'].features['weights'].dtype = 1  # pyrefly: ignore[bad-assignment]
 
     sampling_spec = tfgnn.sampler.SamplingSpec()
     sampling_spec.seed_op.op_name = 'seed'
@@ -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''}
+            keys_to_values={b'a': b''}  # pyrefly: ignore[bad-argument-type]
         )
       else:
         accessor = sampler.InMemIntegerKeyToBytesAccessor(

tensorflow_gnn/experimental/sampler/beam/utils.py

@@ -257,21 +257,21 @@ def parse_tf_example(
     `[c.numpy() for c in [*rt.flat_values, *rt.nested_row_lengths()]]`.
   """
   if spec.HasField('tensor'):
-    spec = spec.tensor
+    spec = spec.tensor  # pyrefly: ignore[bad-assignment]
     return [
         _parse_tf_feature(
             example,
             name,
-            get_np_dtype(spec.dtype),
-            tuple(dim.size for dim in spec.shape.dim),
+            get_np_dtype(spec.dtype),  # pyrefly: ignore[missing-attribute]
+            tuple(dim.size for dim in spec.shape.dim),  # pyrefly: ignore[missing-attribute]
         )
     ]
   elif spec.HasField('ragged_tensor'):
-    spec = spec.ragged_tensor
-    dtype = get_np_dtype(spec.dtype)
-    ragged_rank = spec.ragged_rank
-    row_splits_dtype = get_np_dtype(spec.row_splits_dtype)
-    shape = tuple(dim.size for dim in spec.shape.dim)
+    spec = spec.ragged_tensor  # pyrefly: ignore[bad-assignment]
+    dtype = get_np_dtype(spec.dtype)  # pyrefly: ignore[missing-attribute]
+    ragged_rank = spec.ragged_rank  # pyrefly: ignore[missing-attribute]
+    row_splits_dtype = get_np_dtype(spec.row_splits_dtype)  # pyrefly: ignore[missing-attribute]
+    shape = tuple(dim.size for dim in spec.shape.dim)  # pyrefly: ignore[missing-attribute]
 
     flat_value = _parse_tf_feature(example, name, dtype, shape[ragged_rank:])
 

tensorflow_gnn/experimental/sampler/beam/utils_test.py

@@ -134,9 +134,9 @@ def test_dense_slice(
       limit: int,
       expected: np.ndarray,
   ):
-    value = [value]
-    expected = [expected]
-    actual = utils.ragged_slice(value, start, limit)
+    value = [value]  # pyrefly: ignore[bad-assignment]
+    expected = [expected]  # pyrefly: ignore[bad-assignment]
+    actual = utils.ragged_slice(value, start, limit)  # pyrefly: ignore[bad-argument-type]
     tf.nest.map_structure(self.assertAllEqual, actual, expected)
 
   @parameterized.named_parameters([
@@ -194,9 +194,9 @@ def as_value(r: tf.RaggedTensor) -> List[np.ndarray]:
           lambda t: t.numpy(), [r.flat_values, *r.nested_row_lengths()]
       )
 
-    value = as_value(value)
-    expected = as_value(expected)
-    actual = utils.ragged_slice(value, start, limit)
+    value = as_value(value)  # pyrefly: ignore[bad-assignment]
+    expected = as_value(expected)  # pyrefly: ignore[bad-assignment]
+    actual = utils.ragged_slice(value, start, limit)  # pyrefly: ignore[bad-argument-type]
     tf.nest.map_structure(self.assertAllEqual, actual, expected)
 
 

tensorflow_gnn/graph/adjacency.py

@@ -199,9 +199,9 @@ def _merge_batch_to_components(
     assert isinstance(flat_adj, HyperAdjacency)
 
     def flatten_indices(node_tag_key, index: Field) -> Field:
-      node_set_name = self.spec._metadata[node_tag_key]  # pylint: disable=protected-access
-      return utils.flatten_indices(index, num_edges_per_example,
-                                   num_nodes_per_example[node_set_name])
+      node_set_name = self.spec._metadata[node_tag_key]  # pylint: disable=protected-access  # pyrefly: ignore[unsupported-operation]
+      return utils.flatten_indices(index, num_edges_per_example,  # pyrefly: ignore[bad-argument-type]
+                                   num_nodes_per_example[node_set_name])  # pyrefly: ignore[bad-argument-type, bad-index]
 
     new_data = {
         node_tag_key: flatten_indices(node_tag_key, index)
@@ -324,7 +324,7 @@ def get_index_specs_dict(
 
   def node_set_name(self, node_set_tag: IncidentNodeTag) -> NodeSetName:
     """Returns a node set name for the given node set tag."""
-    return self._metadata[_node_tag_to_index_key(node_set_tag)]
+    return self._metadata[_node_tag_to_index_key(node_set_tag)]  # pyrefly: ignore[bad-return, unsupported-operation]
 
   @property
   def total_size(self) -> Optional[int]:
@@ -477,7 +477,7 @@ class AdjacencySpec(HyperAdjacencySpec):
   """A type spec for `tfgnn.Adjacency`."""
 
   @classmethod
-  def from_incident_node_sets(
+  def from_incident_node_sets(  # pyrefly: ignore[bad-override]
       cls,
       source_node_set: NodeSetName,
       target_node_set: NodeSetName,
@@ -602,25 +602,25 @@ def check_compatibility(tag_0, name_0, index_0, tag_i, name_i, index_i):
 
     raise ValueError(err_message)
 
-  indices = sorted(list(indices.items()), key=lambda i: i[0])
+  indices = sorted(list(indices.items()), key=lambda i: i[0])  # pyrefly: ignore[bad-assignment]
   tag_0, (name_0, index_0) = indices[0]
   check_index(tag_0, name_0, index_0)
-  for tag_i, (name_i, index_i) in indices[1:]:
+  for tag_i, (name_i, index_i) in indices[1:]:  # pyrefly: ignore[bad-index]
     check_index(tag_i, name_i, index_i)
     check_compatibility(tag_0, name_0, index_0, tag_i, name_i, index_i)
 
   # Apply identity operations to all index tensors to ensure that assertions are
   # executed in the graph mode.
   if not assert_ops:
     result = {
-        node_tag: (node_set, index) for node_tag, (node_set, index) in indices
+        node_tag: (node_set, index) for node_tag, (node_set, index) in indices  # pyrefly: ignore[not-iterable]
     }
   else:
     assert allow_tf_assertions
     with tf.control_dependencies(assert_ops):
       result = {
           node_tag: (node_set, tf.identity(index))
-          for node_tag, (node_set, index) in indices
+          for node_tag, (node_set, index) in indices  # pyrefly: ignore[not-iterable]
       }
 
   return result

tensorflow_gnn/graph/batching_utils.py

@@ -145,7 +145,7 @@ def get_next_state(state: _ScanState,
   def exceeds_budget(state: _ScanState,
                      graph_tensor: gt.GraphTensor) -> tf.Tensor:
     budget = _set_min_nodes_per_component(state.budget_left,
-                                          min_nodes_per_component)
+                                          min_nodes_per_component)  # pyrefly: ignore[bad-argument-type]
     within_budget = padding_ops.satisfies_size_constraints(graph_tensor, budget)
     return tf.math.logical_not(within_budget)
 
@@ -631,9 +631,9 @@ def _set_min_nodes_per_component(
 
   return SizeConstraints(
       total_num_components=size_constraints.total_num_components,
-      total_num_nodes=size_constraints.total_num_nodes.copy(),
-      total_num_edges=size_constraints.total_num_edges.copy(),
-      min_nodes_per_component=min_nodes_per_component.copy())
+      total_num_nodes=size_constraints.total_num_nodes.copy(),  # pyrefly: ignore[missing-attribute]
+      total_num_edges=size_constraints.total_num_edges.copy(),  # pyrefly: ignore[missing-attribute]
+      min_nodes_per_component=min_nodes_per_component.copy())  # pyrefly: ignore[missing-attribute]
 
 
 def _make_room_for_padding(
@@ -876,4 +876,4 @@ def restored_generator():
       yield value
 
   return tf.data.Dataset.from_generator(
-      restored_generator, output_signature=relaxed_spec)
+      restored_generator, output_signature=relaxed_spec)  # pyrefly: ignore[bad-argument-type]