diff --git a/tensorflow_gnn/experimental/sampler/beam/accessors_test.py b/tensorflow_gnn/experimental/sampler/beam/accessors_test.py
index 83bb8bab..9193e7e6 100644
--- a/tensorflow_gnn/experimental/sampler/beam/accessors_test.py
+++ b/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),
diff --git a/tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py b/tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py
index a18e1e3b..194e9769 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'
+ 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),
diff --git a/tensorflow_gnn/experimental/sampler/beam/executor_lib.py b/tensorflow_gnn/experimental/sampler/beam/executor_lib.py
index f40038e3..2731dd1a 100644
--- a/tensorflow_gnn/experimental/sampler/beam/executor_lib.py
+++ b/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,
)
diff --git a/tensorflow_gnn/experimental/sampler/beam/executor_lib_test.py b/tensorflow_gnn/experimental/sampler/beam/executor_lib_test.py
index 230a6bd7..abd67679 100644
--- a/tensorflow_gnn/experimental/sampler/beam/executor_lib_test.py
+++ b/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:
# => => identity -> => 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)
diff --git a/tensorflow_gnn/experimental/sampler/beam/sampler.py b/tensorflow_gnn/experimental/sampler/beam/sampler.py
index e9bebecc..bfb4fda7 100644
--- a/tensorflow_gnn/experimental/sampler/beam/sampler.py
+++ b/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
diff --git a/tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py b/tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py
index de44ddc5..607b830c 100644
--- a/tensorflow_gnn/experimental/sampler/beam/subgraph_pipeline_test.py
+++ b/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(
diff --git a/tensorflow_gnn/experimental/sampler/beam/utils.py b/tensorflow_gnn/experimental/sampler/beam/utils.py
index 0df94e67..cd2946d1 100644
--- a/tensorflow_gnn/experimental/sampler/beam/utils.py
+++ b/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:])
diff --git a/tensorflow_gnn/experimental/sampler/beam/utils_test.py b/tensorflow_gnn/experimental/sampler/beam/utils_test.py
index d0f6b0ba..af9bcdee 100644
--- a/tensorflow_gnn/experimental/sampler/beam/utils_test.py
+++ b/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)
diff --git a/tensorflow_gnn/graph/adjacency.py b/tensorflow_gnn/graph/adjacency.py
index 37613ebd..b623c4b4 100644
--- a/tensorflow_gnn/graph/adjacency.py
+++ b/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,10 +602,10 @@ 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)
@@ -613,14 +613,14 @@ def check_compatibility(tag_0, name_0, index_0, tag_i, name_i, index_i):
# 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
diff --git a/tensorflow_gnn/graph/batching_utils.py b/tensorflow_gnn/graph/batching_utils.py
index 45ff28a3..5fea9825 100644
--- a/tensorflow_gnn/graph/batching_utils.py
+++ b/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]
diff --git a/tensorflow_gnn/graph/broadcast_ops.py b/tensorflow_gnn/graph/broadcast_ops.py
index f1d0694f..55629419 100644
--- a/tensorflow_gnn/graph/broadcast_ops.py
+++ b/tensorflow_gnn/graph/broadcast_ops.py
@@ -178,7 +178,7 @@ def _broadcast_context(graph_tensor: GraphTensor,
# TODO(b/184021442): cache result.
return utils.repeat(
context_value,
- node_or_edge_set.sizes,
+ node_or_edge_set.sizes, # pyrefly: ignore[bad-argument-type]
repeats_sum_hint=node_or_edge_set.spec.total_size)
@@ -250,11 +250,11 @@ def broadcast_v2(
for name in edge_set_names]
else:
result = [broadcast_context_to_nodes(graph_tensor, name, **feature_kwargs)
- for name in node_set_names]
+ for name in node_set_names] # pyrefly: ignore[not-iterable]
else:
result = [
broadcast_node_to_edges(graph_tensor, name, from_tag, **feature_kwargs)
- for name in edge_set_names]
+ for name in edge_set_names] # pyrefly: ignore[not-iterable]
if got_sequence_args:
return result
diff --git a/tensorflow_gnn/graph/graph_tensor.py b/tensorflow_gnn/graph/graph_tensor.py
index 9b16e569..894ead3e 100644
--- a/tensorflow_gnn/graph/graph_tensor.py
+++ b/tensorflow_gnn/graph/graph_tensor.py
@@ -513,7 +513,7 @@ def from_field_specs(
shape=sizes_shape,
ragged_rank=shape.rank + 1,
dtype=indices_dtype,
- row_splits_dtype=indicative_feature_spec.row_splits_dtype)
+ row_splits_dtype=indicative_feature_spec.row_splits_dtype) # pyrefly: ignore[missing-attribute]
else:
sizes_spec = tf.TensorSpec(shape=sizes_shape, dtype=indices_dtype)
@@ -1244,9 +1244,9 @@ def edge_set_merge_batch_to_components(edge_set: EdgeSet) -> EdgeSet:
return self.__class__.from_pieces(
context=self.context._merge_batch_to_components(), # pylint: disable=protected-access
node_sets=tf.nest.map_structure(
- lambda n: n._merge_batch_to_components(), self.node_sets.copy()), # pylint: disable=protected-access
+ lambda n: n._merge_batch_to_components(), self.node_sets.copy()), # pylint: disable=protected-access # pyrefly: ignore[missing-attribute]
edge_sets=tf.nest.map_structure(edge_set_merge_batch_to_components,
- self.edge_sets.copy()))
+ self.edge_sets.copy())) # pyrefly: ignore[missing-attribute]
@property
def context(self) -> Context:
@@ -1363,7 +1363,7 @@ def replace_features(
new_context = self.context.replace_features(context)
if node_sets is None:
- new_node_sets = self.node_sets.copy()
+ new_node_sets = self.node_sets.copy() # pyrefly: ignore[missing-attribute]
else:
not_present = set(node_sets.keys()) - set(self.node_sets.keys())
if not_present:
@@ -1376,7 +1376,7 @@ def replace_features(
}
if edge_sets is None:
- new_edge_sets = self.edge_sets.copy()
+ new_edge_sets = self.edge_sets.copy() # pyrefly: ignore[missing-attribute]
else:
not_present = set(edge_sets.keys()) - set(self.edge_sets.keys())
if not_present:
@@ -1778,7 +1778,7 @@ def _fields_and_size_from_fieldorfields(
"""Returns a mapping from a default feature name if needed."""
if isinstance(features, collections.abc.Mapping):
num_entities = tf.stack(
- [utils.outer_dimension_size(_get_indicative_feature(features))])
+ [utils.outer_dimension_size(_get_indicative_feature(features))]) # pyrefly: ignore[bad-argument-type]
elif features is None:
features, num_entities = {}, None
else:
@@ -1833,11 +1833,11 @@ def homogeneous(
raise ValueError('source and target must be rank-1 dense tensors')
node_features, num_nodes = _fields_and_size_from_fieldorfields(
- node_features, HIDDEN_STATE)
+ node_features, HIDDEN_STATE) # pyrefly: ignore[bad-argument-type]
edge_features, _ = _fields_and_size_from_fieldorfields(
- edge_features, HIDDEN_STATE)
+ edge_features, HIDDEN_STATE) # pyrefly: ignore[bad-argument-type]
context_features, _ = _fields_and_size_from_fieldorfields(
- context_features, HIDDEN_STATE)
+ context_features, HIDDEN_STATE) # pyrefly: ignore[bad-argument-type]
num_edges = tf.shape(source)
node_sizes = (num_nodes if node_set_sizes is None else node_set_sizes)
diff --git a/tensorflow_gnn/graph/graph_tensor_ops.py b/tensorflow_gnn/graph/graph_tensor_ops.py
index 4e6d0ba0..2d1f68c4 100644
--- a/tensorflow_gnn/graph/graph_tensor_ops.py
+++ b/tensorflow_gnn/graph/graph_tensor_ops.py
@@ -100,9 +100,9 @@ def add_self_loops(
# |E1| |N1| |E2| |N2|
segment_indicator = utils.repeat(
tf.range(tf.shape(alternate_sizes)[0], dtype=tf.int32), alternate_sizes,
- repeats_sum_hint=tf.get_static_value(num_nodes + num_edges))
+ repeats_sum_hint=tf.get_static_value(num_nodes + num_edges)) # pyrefly: ignore[unsupported-operation]
- node_indicator = segment_indicator % 2 # Marks odd (i.e. node positions)
+ node_indicator = segment_indicator % 2 # Marks odd (i.e. node positions) # pyrefly: ignore[unsupported-operation]
edge_indicator = 1 - node_indicator # Marks even (i.e. edge positions)
# [0, 1, 2,.., x, x, ..., |E1|, |E1|+1,.., x, x, x, ...]; "x" = dont care.
@@ -148,7 +148,7 @@ def add_self_loops(
self_loop_edge_feature = utils.repeat(
tf.expand_dims(self_loop_edge_feature, axis=0),
tf.expand_dims(num_nodes, axis=0),
- repeats_sum_hint=tf.get_static_value(num_nodes + 0))
+ repeats_sum_hint=tf.get_static_value(num_nodes + 0)) # pyrefly: ignore[unsupported-operation]
# Transposing twice so that we get broadcasting for free (instead of
# reshaping, adding 1's on some axis dimensions).
# TODO(b/309749041): What if there are no existing feature values?
@@ -293,7 +293,7 @@ def mask_edges(
masked_info_features[feature_name] = masked_info_feature
component_ids = utils.row_lengths_to_row_ids(
- edge_set.sizes, sum_row_lengths_hint=edge_set.spec.total_size)
+ edge_set.sizes, sum_row_lengths_hint=edge_set.spec.total_size) # pyrefly: ignore[bad-argument-type]
num_remaining_edges = tf.math.unsorted_segment_sum(
tf.cast(boolean_edge_mask, edge_set.sizes.dtype), component_ids,
edge_set.num_components)
@@ -592,7 +592,7 @@ def index_shuffle_singleton(node_set: gt.NodeSet) -> tf.Tensor:
def index_shuffle_generic(node_set: gt.NodeSet) -> tf.Tensor:
row_ids = utils.row_lengths_to_row_ids(
- node_set.sizes, sum_row_lengths_hint=node_set.spec.total_size)
+ node_set.sizes, sum_row_lengths_hint=node_set.spec.total_size) # pyrefly: ignore[bad-argument-type]
return utils.segment_random_index_shuffle(segment_ids=row_ids, seed=seed)
if graph_tensor.spec.total_num_components == 1:
@@ -902,7 +902,7 @@ def _get_line_graph_edges(
# e.g. [0 0 1 1 2 2 2 3 3 3]
idx_line_source = utils.repeat(
edge_idx_sorted_source,
- utils.repeat(num_neighbors_target, num_neighbors_source),
+ utils.repeat(num_neighbors_target, num_neighbors_source), # pyrefly: ignore[bad-argument-type]
)
# Repeat target line idx according to the number of neighbors per node
diff --git a/tensorflow_gnn/graph/graph_tensor_random.py b/tensorflow_gnn/graph/graph_tensor_random.py
index e109a982..c6f6efff 100644
--- a/tensorflow_gnn/graph/graph_tensor_random.py
+++ b/tensorflow_gnn/graph/graph_tensor_random.py
@@ -88,7 +88,7 @@ def random_ragged_tensor(
sample_values_tensor = tf.convert_to_tensor(sample_values, dtype=dtype)
flat_values = tf.gather(sample_values_tensor, indices)
else:
- flat_values = typed_random_values(size, dtype)
+ flat_values = typed_random_values(size, dtype) # pyrefly: ignore[bad-argument-type]
# Now, build up the ragged tensor inside out.
#
@@ -222,7 +222,7 @@ def _gen_features(
# Create random context features.
context = gt.Context.from_fields(
features=_gen_features(
- gc.CONTEXT, None, spec.context_spec.features_spec, num_components
+ gc.CONTEXT, None, spec.context_spec.features_spec, num_components # pyrefly: ignore[bad-argument-type]
)
)
@@ -231,7 +231,7 @@ def _gen_features(
for set_name, node_set_spec in spec.node_sets_spec.items():
min_nodes, max_nodes = row_lengths_range
sizes = _random_sizes(
- num_components, min_nodes, max_nodes, spec.indices_dtype
+ num_components, min_nodes, max_nodes, spec.indices_dtype # pyrefly: ignore[bad-argument-type]
)
node_sets[set_name] = gt.NodeSet.from_fields(
sizes=sizes,
@@ -260,7 +260,7 @@ def _gen_features(
min_edges = tf.cast(sum_sizes / 1.5, spec.indices_dtype)
max_edges = tf.cast(sum_sizes * 2.25, spec.indices_dtype)
sizes = _random_sizes(
- num_components, min_edges, max_edges, spec.indices_dtype
+ num_components, min_edges, max_edges, spec.indices_dtype # pyrefly: ignore[bad-argument-type]
)
# Randomly generate the actual node indices.
@@ -334,7 +334,7 @@ def _random_sizes(
) -> tf.Tensor:
"""Random sizes with constraints on the number of items in each component."""
minval = tf.convert_to_tensor(num_items_min, dtype)
- length = tf.convert_to_tensor(num_items_max - num_items_min, dtype)
+ length = tf.convert_to_tensor(num_items_max - num_items_min, dtype) # pyrefly: ignore[unsupported-operation]
alpha = tf.random.uniform([num_components], dtype=tf.float64)
return minval + tf.cast(alpha * tf.cast(length, tf.float64), dtype)
diff --git a/tensorflow_gnn/graph/padding_ops.py b/tensorflow_gnn/graph/padding_ops.py
index 5ea02fe7..3c1018d3 100644
--- a/tensorflow_gnn/graph/padding_ops.py
+++ b/tensorflow_gnn/graph/padding_ops.py
@@ -181,7 +181,7 @@ def get_min_max_fake_nodes_indices(
tf.ones([total_num_components], dtype=tf.bool),
tf.zeros([num_padded], dtype=tf.bool)
],
- axis=0), target_total_num_components)
+ axis=0), target_total_num_components) # pyrefly: ignore[bad-argument-type]
return padded_graph_tensor, cast(tf.Tensor, padding_mask)
@@ -269,7 +269,7 @@ def _(context: gt.Context, *,
"""Pads graph context to the target number of graph components."""
diff = tf.ones(
- shape=[target_total_num_components - context.total_num_components],
+ shape=[target_total_num_components - context.total_num_components], # pyrefly: ignore[unsupported-operation]
dtype=context.spec.sizes_spec.dtype)
sizes = tf.concat([context.sizes, diff], axis=0)
sizes = tensor_utils.ensure_static_nrows(
@@ -415,7 +415,7 @@ def _pad_adjacency_index_with_linspace(index: const.Field, target_size: int,
target_size, dtype=index.dtype) - tf.size(index, index.dtype)
diff = tf.linspace(
start=tf.cast(min_index, tf.float32),
- stop=tf.cast(max_index + 1, tf.float32),
+ stop=tf.cast(max_index + 1, tf.float32), # pyrefly: ignore[unsupported-operation]
num=diff_size)
diff = tf.cast(diff, index.dtype)
diff = tf.clip_by_value(diff, min_index, max_index)
@@ -522,14 +522,14 @@ def _satisfies_size_constraints_internal(
# np.ndarray and tf.Tensor they could be evaluated statically on in the
# runtime depending on its arguments.
total_num_components = graph_tensor.total_num_components
- could_add_new_component = _fold_constants(lambda x, y: x < y,
+ could_add_new_component = _fold_constants(lambda x, y: x < y, # pyrefly: ignore[bad-argument-type, unsupported-operation]
total_num_components,
- total_sizes.total_num_components)
- num_fake_components = total_sizes.total_num_components - graph_tensor.total_num_components
+ total_sizes.total_num_components) # pyrefly: ignore[bad-argument-type]
+ num_fake_components = total_sizes.total_num_components - graph_tensor.total_num_components # pyrefly: ignore[unsupported-operation]
assert_ops = [
check_fn(
_fold_constants(
- lambda x, y: x <= y, total_num_components,
+ lambda x, y: x <= y, total_num_components, # pyrefly: ignore[bad-argument-type, unsupported-operation]
tf.convert_to_tensor(
total_sizes.total_num_components,
dtype=total_num_components.dtype)),
@@ -581,21 +581,21 @@ def _check_sizes(entity_type: str, entity_name: str,
# because the weaker case may support constant folding.
assert_ops.append(
check_fn(
- _fold_constants(lambda x, y: x <= y, total_size,
+ _fold_constants(lambda x, y: x <= y, total_size, # pyrefly: ignore[bad-argument-type, unsupported-operation]
target_total_size),
overflow_msg))
assert_ops.append(
check_fn(
- _fold_constants(lambda x, y: x <= y, padded_size,
+ _fold_constants(lambda x, y: x <= y, padded_size, # pyrefly: ignore[bad-argument-type, unsupported-operation]
target_total_size),
overflow_msg))
assert_ops.append(
check_fn(
_fold_constants(
- lambda x, y: x | y, could_add_new_component,
- _fold_constants(lambda x, y: x == y, padded_size,
+ lambda x, y: x | y, could_add_new_component, # pyrefly: ignore[unsupported-operation]
+ _fold_constants(lambda x, y: x == y, padded_size, # pyrefly: ignore[bad-argument-type]
target_total_size)),
(f'Could not pad <{entity_name}> {entity_type}. To do this, at'
' least one graph component must be added to the input graph.'
@@ -611,28 +611,28 @@ def _check_sizes(entity_type: str, entity_name: str,
_check_sizes(
'nodes',
name,
- item.sizes,
+ item.sizes, # pyrefly: ignore[bad-argument-type]
total_size,
- target_total_size,
- min_entities_per_component=min_nodes_per_component.get(name, 0))
+ target_total_size, # pyrefly: ignore[bad-argument-type]
+ min_entities_per_component=min_nodes_per_component.get(name, 0)) # pyrefly: ignore[bad-argument-type]
for name, item in graph_tensor.edge_sets.items():
total_size = item.total_size
target_total_size = total_sizes.total_num_edges.get(name, None)
- _check_sizes('edges', name, item.sizes, total_size, target_total_size,
+ _check_sizes('edges', name, item.sizes, total_size, target_total_size, # pyrefly: ignore[bad-argument-type]
min_entities_per_component=0)
assert target_total_size is not None
- has_all_edges = _fold_constants(lambda x, y: x == y, total_size,
- target_total_size)
+ has_all_edges = _fold_constants(lambda x, y: x == y, total_size, # pyrefly: ignore[bad-argument-type]
+ target_total_size) # pyrefly: ignore[bad-argument-type]
indices = item.adjacency.get_indices_dict()
for _, (incident_node_set_name, _) in indices.items():
permits_new_incident_nodes = _fold_constants(
- lambda x, y: x < y, total_num_nodes[incident_node_set_name],
- total_sizes.total_num_nodes[incident_node_set_name])
+ lambda x, y: x < y, total_num_nodes[incident_node_set_name], # pyrefly: ignore[bad-argument-type, unsupported-operation]
+ total_sizes.total_num_nodes[incident_node_set_name]) # pyrefly: ignore[bad-argument-type]
assert_ops.append(
check_fn(
- _fold_constants(lambda x, y: x | y, has_all_edges,
+ _fold_constants(lambda x, y: x | y, has_all_edges, # pyrefly: ignore[unsupported-operation]
permits_new_incident_nodes),
('Could not create fake incident edges for the node set'
f' {incident_node_set_name}. This could happen when the'
diff --git a/tensorflow_gnn/graph/pool_ops.py b/tensorflow_gnn/graph/pool_ops.py
index 78c9ed84..02c02387 100644
--- a/tensorflow_gnn/graph/pool_ops.py
+++ b/tensorflow_gnn/graph/pool_ops.py
@@ -322,7 +322,7 @@ def pool_v2(
else:
msg_lines.extend([
f" node set '{name}' has feature shape {shape.as_list()}"
- for name, shape in zip(node_set_names, feature_shapes)])
+ for name, shape in zip(node_set_names, feature_shapes)]) # pyrefly: ignore[bad-argument-type]
raise ValueError("\n".join(msg_lines))
return _pool_internal(
@@ -477,15 +477,15 @@ def get_pool_args_as_sequences(
elif feature_value is not None:
feature_values = [feature_value]
elif edge_set_names is not None:
- feature_values = [graph.edge_sets[edge_set_name][feature_name]
+ feature_values = [graph.edge_sets[edge_set_name][feature_name] # pyrefly: ignore[bad-index]
for edge_set_name in edge_set_names]
else:
- feature_values = [graph.node_sets[node_set_name][feature_name]
- for node_set_name in node_set_names]
+ feature_values = [graph.node_sets[node_set_name][feature_name] # pyrefly: ignore[bad-index]
+ for node_set_name in node_set_names] # pyrefly: ignore[not-iterable]
if edge_set_names is not None:
_check_same_length("edge_set_names", edge_set_names, feature_values)
else:
- _check_same_length("node_set_names", node_set_names, feature_values)
+ _check_same_length("node_set_names", node_set_names, feature_values) # pyrefly: ignore[bad-argument-type]
return edge_set_names, node_set_names, feature_values, got_sequence_args
@@ -537,24 +537,24 @@ def reduce(
if edge_set_name is not None:
node_or_edge_set = graph.edge_sets[edge_set_name]
else:
- node_or_edge_set = graph.node_sets[node_set_name]
+ node_or_edge_set = graph.node_sets[node_set_name] # pyrefly: ignore[bad-index]
sizes = node_or_edge_set.sizes
return self.unsorted_segment_op(
feature_value,
utils.row_lengths_to_row_ids(
- sizes, sum_row_lengths_hint=node_or_edge_set.spec.total_size),
- utils.outer_dimension_size(sizes))
+ sizes, sum_row_lengths_hint=node_or_edge_set.spec.total_size), # pyrefly: ignore[bad-argument-type]
+ utils.outer_dimension_size(sizes)) # pyrefly: ignore[bad-argument-type]
# Pooling from edges to node.
- adjacency = graph.edge_sets[edge_set_name].adjacency
+ adjacency = graph.edge_sets[edge_set_name].adjacency # pyrefly: ignore[bad-index]
if isinstance(adjacency, (kt.HyperAdjacencyKerasTensor, # TODO(b/283404258)
adj.HyperAdjacency)):
- node_set = graph.node_sets[adjacency.node_set_name(to_tag)]
+ node_set = graph.node_sets[adjacency.node_set_name(to_tag)] # pyrefly: ignore[bad-argument-type]
total_node_count = node_set.spec.total_size
if total_node_count is None:
total_node_count = node_set.total_size
return self.unsorted_segment_op(feature_value,
- adjacency[to_tag],
+ adjacency[to_tag], # pyrefly: ignore[bad-argument-type, bad-index]
total_node_count)
else:
raise ValueError(f"Edge set '{edge_set_name}' has unknown "
diff --git a/tensorflow_gnn/graph/preprocessing_common.py b/tensorflow_gnn/graph/preprocessing_common.py
index d50e861b..f17b51ee 100644
--- a/tensorflow_gnn/graph/preprocessing_common.py
+++ b/tensorflow_gnn/graph/preprocessing_common.py
@@ -57,7 +57,7 @@ def compute_basic_stats(dataset: tf.data.Dataset) -> BasicStats:
def reduce_fn(old_state: Tuple[tf.Tensor, Any],
new_element) -> Tuple[tf.Tensor, Any]:
old_count, old_stats = old_state
- new_count = old_count + 1
+ new_count = old_count + 1 # pyrefly: ignore[unsupported-operation]
alpha = 1.0 / tf.cast(new_count, tf.float64)
new_stats = BasicStats(
minimum=tf.nest.map_structure(tf.minimum, old_stats.minimum,
@@ -71,7 +71,7 @@ def reduce_fn(old_state: Tuple[tf.Tensor, Any],
lambda s, e: s + (tf.cast(e, tf.float64) - s) * alpha,
old_stats.mean, new_element),
)
- return (new_count, new_stats)
+ return (new_count, new_stats) # pyrefly: ignore[bad-return]
def cast_to_float32_or_float64(value_spec, value):
if value_spec.dtype == value.dtype:
@@ -154,7 +154,7 @@ def scan_fn(
pred = predicate(value)
ema_update = tf.cast(tf.logical_not(pred), old_ema.dtype)
- in_count = old_in_count + 1
+ in_count = old_in_count + 1 # pyrefly: ignore[unsupported-operation]
out_count = old_out_count + tf.cast(pred, old_out_count.dtype)
# Exponential moving average (EMA) is updated according to the rule
@@ -178,7 +178,7 @@ def scan_fn(
ema = tf.cond(in_count % summary_steps == 0,
lambda: report(ema, step=report_step), lambda: ema)
- return (in_count, out_count, ema), (pred, value)
+ return (in_count, out_count, ema), (pred, value) # pyrefly: ignore[bad-return]
state0 = (tf.constant(0, tf.int64), tf.constant(0, tf.int64),
tf.constant(0., tf.float64))
diff --git a/tensorflow_gnn/graph/readout.py b/tensorflow_gnn/graph/readout.py
index d98b6586..f3d1e3ec 100644
--- a/tensorflow_gnn/graph/readout.py
+++ b/tensorflow_gnn/graph/readout.py
@@ -422,7 +422,7 @@ def context_readout_into_feature(
"Pass tfgnn.context_readout_into_feature(..., overwrite=True) "
"to discard the old value."
)
- readout_features[new_feature_name] = input_feature
+ readout_features[new_feature_name] = input_feature # pyrefly: ignore[unsupported-operation]
return graph.replace_features(node_sets={readout_node_set: readout_features},
context=context_features)
diff --git a/tensorflow_gnn/graph/schema_utils.py b/tensorflow_gnn/graph/schema_utils.py
index 4ac67b6f..56d12eb2 100644
--- a/tensorflow_gnn/graph/schema_utils.py
+++ b/tensorflow_gnn/graph/schema_utils.py
@@ -36,7 +36,7 @@ def parse_schema(schema_text: Union[bytes, str]) -> schema_pb2.GraphSchema:
Returns:
A `GraphSchema` instance.
"""
- return text_format.Parse(schema_text, schema_pb2.GraphSchema())
+ return text_format.Parse(schema_text, schema_pb2.GraphSchema()) # pyrefly: ignore[bad-specialization]
def read_schema(filename: str) -> schema_pb2.GraphSchema:
diff --git a/tensorflow_gnn/graph/schema_validation.py b/tensorflow_gnn/graph/schema_validation.py
index 6b0b165f..a2f8ace6 100644
--- a/tensorflow_gnn/graph/schema_validation.py
+++ b/tensorflow_gnn/graph/schema_validation.py
@@ -322,7 +322,7 @@ def assert_constraints(graph: gt.GraphTensor) -> tf.Operation:
def _assert_constraints_feature_shape_prefix(
graph: gt.GraphTensor) -> tf.Operation:
"""Validates the number of nodes or edges of feature tensors."""
- with tf.name_scope("constraints_feature_shape_prefix"):
+ with tf.name_scope("constraints_feature_shape_prefix"): # pyrefly: ignore[bad-instantiation]
checks = []
for set_type, set_dict in [("node", graph.node_sets),
("edge", graph.edge_sets)]:
@@ -354,7 +354,7 @@ def _assert_constraints_edge_indices_range(
graph: gt.GraphTensor) -> tf.Operation:
"""Validates that edge indices are within the bounds of node set sizes."""
- with tf.name_scope("constraints_edge_indices_range"):
+ with tf.name_scope("constraints_edge_indices_range"): # pyrefly: ignore[bad-instantiation]
checks = []
for set_name, edge_set in graph.edge_sets.items():
adjacency = edge_set.adjacency
@@ -390,7 +390,7 @@ def _assert_constraints_edge_indices_range(
def _assert_constraints_edge_shapes(graph: gt.GraphTensor) -> tf.Operation:
"""Validates edge shapes and that they contain a scalar index per node."""
- with tf.name_scope("constraints_edge_indices_range"):
+ with tf.name_scope("constraints_edge_indices_range"): # pyrefly: ignore[bad-instantiation]
checks = []
for set_name, edge_set in graph.edge_sets.items():
adjacency = edge_set.adjacency
diff --git a/tensorflow_gnn/graph/tag_utils.py b/tensorflow_gnn/graph/tag_utils.py
index 72afd27e..840fbccf 100644
--- a/tensorflow_gnn/graph/tag_utils.py
+++ b/tensorflow_gnn/graph/tag_utils.py
@@ -103,7 +103,7 @@ def get_edge_or_node_set_name_args_for_tag(
if tag != const.CONTEXT:
incident_node_set_names = {
spec.edge_sets_spec[e].adjacency_spec.node_set_name(tag)
- for e in edge_set_names}
+ for e in edge_set_names} # pyrefly: ignore[not-iterable]
if len(incident_node_set_names) > 1:
raise ValueError(
f"{function_name} requires the same endpoint for all named edge sets "
diff --git a/tensorflow_gnn/graph/tensor_utils.py b/tensorflow_gnn/graph/tensor_utils.py
index 7d9ab285..9e62865c 100644
--- a/tensorflow_gnn/graph/tensor_utils.py
+++ b/tensorflow_gnn/graph/tensor_utils.py
@@ -39,7 +39,7 @@ def outer_dimension_size(value: Value) -> Union[int, tf.Tensor]:
return outer_dimension_size(value.row_lengths())
if is_dense_tensor(value):
- return dims_list(value)[0]
+ return dims_list(value)[0] # pyrefly: ignore[bad-argument-type]
raise ValueError(f'Unsupported type {type(value).__name__}')
@@ -168,7 +168,7 @@ def flatten_indices(indices: tf.Tensor, indices_row_lengths: tf.Tensor,
_assert_rank1_int(values_row_lengths, 'values_row_lengths')
indices_total = outer_dimension_size(indices)
- indices_row_ids = row_lengths_to_row_ids(indices_row_lengths, indices_total)
+ indices_row_ids = row_lengths_to_row_ids(indices_row_lengths, indices_total) # pyrefly: ignore[bad-argument-type]
values_row_starts = tf.math.cumsum(values_row_lengths, axis=0, exclusive=True)
offsets = tf.gather(values_row_starts, indices_row_ids)
diff --git a/tensorflow_gnn/keras/layers/convolution_base.py b/tensorflow_gnn/keras/layers/convolution_base.py
index 2352f66a..0cc7c83f 100644
--- a/tensorflow_gnn/keras/layers/convolution_base.py
+++ b/tensorflow_gnn/keras/layers/convolution_base.py
@@ -277,7 +277,7 @@ def call(self, graph: gt.GraphTensor, *,
# Pooling from NodeSet to Context, no EdgeSet involved.
name_kwarg = dict(node_set_name=node_set_name)
edge_set = None
- sender_node_set = graph.node_sets[node_set_name]
+ sender_node_set = graph.node_sets[node_set_name] # pyrefly: ignore[bad-index]
# Values are computed per sender node, no need to broadcast
broadcast_from_sender_node = lambda feature_value: feature_value
receiver_piece = graph.context
@@ -319,19 +319,19 @@ def bind_receiver_args(fn):
if self._receiver_feature is not None:
receiver_input = receiver_piece[self._receiver_feature]
if None not in [sender_node_set, self._sender_node_feature]:
- sender_node_input = sender_node_set[self._sender_node_feature]
+ sender_node_input = sender_node_set[self._sender_node_feature] # pyrefly: ignore[bad-index, unsupported-operation]
if None not in [edge_set, self._sender_edge_feature]:
- sender_edge_input = edge_set[self._sender_edge_feature]
+ sender_edge_input = edge_set[self._sender_edge_feature] # pyrefly: ignore[bad-index, unsupported-operation]
return self.convolve(
- sender_node_input=sender_node_input,
- sender_edge_input=sender_edge_input,
- receiver_input=receiver_input,
- broadcast_from_sender_node=broadcast_from_sender_node,
+ sender_node_input=sender_node_input, # pyrefly: ignore[bad-argument-type]
+ sender_edge_input=sender_edge_input, # pyrefly: ignore[bad-argument-type]
+ receiver_input=receiver_input, # pyrefly: ignore[bad-argument-type]
+ broadcast_from_sender_node=broadcast_from_sender_node, # pyrefly: ignore[bad-argument-type]
broadcast_from_receiver=broadcast_from_receiver,
pool_to_receiver=pool_to_receiver,
**extra_receiver_ops_kwarg,
- training=training)
+ training=training) # pyrefly: ignore[bad-argument-type]
@abc.abstractmethod
def convolve(self, *,
diff --git a/tensorflow_gnn/keras/layers/convolutions.py b/tensorflow_gnn/keras/layers/convolutions.py
index a4a740a2..24264d2c 100644
--- a/tensorflow_gnn/keras/layers/convolutions.py
+++ b/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)
+ messages = self._message_fn(combined_input) # pyrefly: ignore[not-callable]
pooled_messages = pool_to_receiver(messages, reduce_type=self._reduce_type)
return pooled_messages
diff --git a/tensorflow_gnn/keras/layers/graph_update.py b/tensorflow_gnn/keras/layers/graph_update.py
index bbf7ecce..00d59ba4 100644
--- a/tensorflow_gnn/keras/layers/graph_update.py
+++ b/tensorflow_gnn/keras/layers/graph_update.py
@@ -207,8 +207,8 @@ def get_config(self):
"to trigger deferred initialization before it can be saved.")
return dict(
# Sublayers need to be top-level objects in the config (b/209560043).
- **du.with_key_prefix(self._edge_set_updates, "edge_sets/"),
- **du.with_key_prefix(self._node_set_updates, "node_sets/"),
+ **du.with_key_prefix(self._edge_set_updates, "edge_sets/"), # pyrefly: ignore[bad-argument-type]
+ **du.with_key_prefix(self._node_set_updates, "node_sets/"), # pyrefly: ignore[bad-argument-type]
context=self._context_update,
**super().get_config())
@@ -221,7 +221,7 @@ def from_config(cls, config):
def call(self, graph: gt.GraphTensor) -> gt.GraphTensor:
if not self._is_initialized:
with tf.init_scope():
- self._init_from_updates(**self._deferred_init_callback(graph.spec))
+ self._init_from_updates(**self._deferred_init_callback(graph.spec)) # pyrefly: ignore[not-callable]
self._deferred_init_callback = None # Enable garbage collection.
assert self._is_initialized
diff --git a/tensorflow_gnn/keras/layers/item_dropout.py b/tensorflow_gnn/keras/layers/item_dropout.py
index 1068e613..f9c9bec8 100644
--- a/tensorflow_gnn/keras/layers/item_dropout.py
+++ b/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)
+ return self._dropout(inputs) # pyrefly: ignore[not-callable]
diff --git a/tensorflow_gnn/keras/layers/map_features.py b/tensorflow_gnn/keras/layers/map_features.py
index 61eaced5..ad4d6aba 100644
--- a/tensorflow_gnn/keras/layers/map_features.py
+++ b/tensorflow_gnn/keras/layers/map_features.py
@@ -211,9 +211,9 @@ def __init__(self,
self._context_fn = None
self._node_sets_fn = None
self._edge_sets_fn = None
- self._context_model = context_model
- self._node_set_models = node_set_models
- self._edge_set_models = edge_set_models
+ self._context_model = context_model # pyrefly: ignore[unbound-name]
+ self._node_set_models = node_set_models # pyrefly: ignore[unbound-name]
+ self._edge_set_models = edge_set_models # pyrefly: ignore[unbound-name]
self._is_initialized = True
self._allowed_aux_node_sets_pattern = allowed_aux_node_sets_pattern
self._allowed_aux_edge_sets_pattern = allowed_aux_edge_sets_pattern
@@ -225,8 +225,8 @@ def get_config(self):
return dict(
context_model=self._context_model,
# Sublayers need to be top-level objects in the config (b/209560043).
- **du.with_key_prefix(self._node_set_models, "node_set_models/"),
- **du.with_key_prefix(self._edge_set_models, "edge_set_models/"),
+ **du.with_key_prefix(self._node_set_models, "node_set_models/"), # pyrefly: ignore[bad-argument-type]
+ **du.with_key_prefix(self._edge_set_models, "edge_set_models/"), # pyrefly: ignore[bad-argument-type]
allowed_aux_node_sets_pattern=self._allowed_aux_node_sets_pattern,
allowed_aux_edge_sets_pattern=self._allowed_aux_edge_sets_pattern,
**super().get_config())
@@ -274,7 +274,7 @@ def call(self, graph: gt.GraphTensor) -> gt.GraphTensor:
node_set_features = {}
for node_set_name, node_set in graph.node_sets.items():
try:
- model = self._node_set_models[node_set_name]
+ model = self._node_set_models[node_set_name] # pyrefly: ignore[unsupported-operation]
if model is None: continue # Was explicitly ignored in initialization.
except KeyError as e:
if self._ignore_node_set(node_set_name):
@@ -287,7 +287,7 @@ def call(self, graph: gt.GraphTensor) -> gt.GraphTensor:
edge_set_features = {}
for edge_set_name, edge_set in graph.edge_sets.items():
try:
- model = self._edge_set_models[edge_set_name]
+ model = self._edge_set_models[edge_set_name] # pyrefly: ignore[unsupported-operation]
if model is None: continue # Was explicitly ignored in initialization.
except KeyError as e:
if self._ignore_edge_set(edge_set_name):
diff --git a/tensorflow_gnn/keras/layers/next_state.py b/tensorflow_gnn/keras/layers/next_state.py
index f066a558..1c4a04e5 100644
--- a/tensorflow_gnn/keras/layers/next_state.py
+++ b/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)
+ net = self._transformation(net) # pyrefly: ignore[not-callable]
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)
+ net = self._residual_block(net) # pyrefly: ignore[not-callable]
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)
+ net = self._activation(net) # pyrefly: ignore[not-callable]
return net