From 304a611f28482c5373092f1507bdd79d7678dfbb Mon Sep 17 00:00:00 2001
From: Oleh Prypin <oprypin@google.com>
Date: Wed, 9 Sep 2026 11:53:19 -0700
Subject: [PATCH] Suppress new Pyrefly findings before upgrading to v1.2.0

PiperOrigin-RevId: 978671276
---
 tensorflow_gnn/data/unigraph.py                           | 4 ++--
 .../experimental/sampler/beam/edge_samplers_test.py       | 8 ++++----
 .../experimental/sampler/beam/subgraph_pipeline_test.py   | 2 +-
 .../experimental/sampler/subgraph_pipeline_test.py        | 2 +-
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow_gnn/data/unigraph.py b/tensorflow_gnn/data/unigraph.py
index fb17bb88..269097ff 100644
--- a/tensorflow_gnn/data/unigraph.py
+++ b/tensorflow_gnn/data/unigraph.py
@@ -1067,8 +1067,8 @@ def iter_graph_via_schema(schema: tfgnn.GraphSchema) -> Dict[
               Iterable[Tuple[bytes, bytes, Example]],
               Iterable[Tuple[bytes, Example]]]]]:
     return {  # pyrefly: ignore[bad-return]
-        tfgnn.NODES: DictStreams.iter_nodes_via_schema(schema),
-        tfgnn.EDGES: DictStreams.iter_edges_via_schema(schema),
+        tfgnn.NODES: DictStreams.iter_nodes_via_schema(schema),  # pyrefly: ignore[bad-assignment]
+        tfgnn.EDGES: DictStreams.iter_edges_via_schema(schema),  # pyrefly: ignore[bad-assignment]
     }
 
 
diff --git a/tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py b/tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py
index 194e9769..41879e18 100644
--- a/tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py
+++ b/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]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, 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]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, 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]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, 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]
+                keys_to_values={b'?': b''}, name='edges'  # pyrefly: ignore[bad-argument-type, bad-assignment]
             ),
             features_spec={
                 'neighbors': tf.TensorSpec([None], tf.int64),
diff --git a/tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py b/tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py
index 607b830c..74779ce5 100644
--- a/tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py
+++ b/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]
+            keys_to_values={b'a': b''}  # pyrefly: ignore[bad-argument-type, bad-assignment]
         )
       else:
         accessor = sampler.InMemIntegerKeyToBytesAccessor(
diff --git a/tensorflow_gnn/experimental/sampler/subgraph_pipeline_test.py b/tensorflow_gnn/experimental/sampler/subgraph_pipeline_test.py
index c6abc69a..542a52e1 100644
--- a/tensorflow_gnn/experimental/sampler/subgraph_pipeline_test.py
+++ b/tensorflow_gnn/experimental/sampler/subgraph_pipeline_test.py
@@ -398,7 +398,7 @@ def edge_sampler_factory(sampling_op):
       self.assertEqual(sampling_op.edge_set_name, 'a->b')
       if ids_dtype == tf.string:
         accessor = core.InMemStringKeyToBytesAccessor(
-            keys_to_values={b'a': b''}  # pyrefly: ignore[bad-argument-type]
+            keys_to_values={b'a': b''}  # pyrefly: ignore[bad-argument-type, bad-assignment]
         )
       else:
         accessor = core.InMemIntegerKeyToBytesAccessor(keys_to_values={0: b''})