EdgeSet.md

tfgnn.EdgeSet

[TOC]

View source on GitHub

A composite tensor for edge set features, size and adjacency information.

tfgnn.EdgeSet(
    data: Data, spec: 'GraphPieceSpecBase', validate: bool = False
)

Each edge set contains edges as its items that connect nodes from particular node sets. The information which edges connect which nodes is encapsulated in the EdgeSet.adjacency composite tensor (see adjacency.py).

All edges in a edge set have the same features, identified by a string key. Each feature is stored as one tensor and has shape [*graph_shape, num_edges, *feature_shape]. The num_edges is a number of edges in a graph (could be ragged). The feature_shape is a shape of the feature value for each edge. EdgeSet supports both fixed-size and variable-size features. The fixed-size features must have fully defined feature_shape. They are stored as tf.Tensor if num_edges is fixed-size or graph_shape.rank = 0. Variable-size edge features are always stored as tf.RaggedTensor.

Note that edge set features are indexed without regard to graph components. The information which edge belong to which graph component is contained in the .sizes tensor which defines the number of edges in each graph component.

Args
`data`	Nest of Field or subclasses of GraphPieceBase.
`spec`	A subclass of GraphPieceSpecBase with a `_data_spec` that matches `data`.
`validate`	if set, checks that data and spec are aligned, compatible and supported.

Attributes
`adjacency`	The information which edges connect which nodes (see tfgnn.Adjacency).
`features`	A read-only mapping of feature name to feature specs.
`indices_dtype`	The integer type to represent ragged splits.
`num_components`	The number of graph components for each graph.
`rank`	The rank of this Tensor. Guaranteed not to be `None`.
`shape`	A possibly-partial shape specification for this Tensor. The returned TensorShape is guaranteed to have a known rank, but the individual dimension sizes may be unknown.
`sizes`	The number of items in each graph component.
`spec`	The public type specification of this tensor.
`total_num_components`	The total number of graph components.
`total_size`	The total number of items.

Methods

from_fields

View source

@classmethod
from_fields(
    *, features: Optional[Fields] = None, sizes: Field, adjacency: Adjacency
) -> 'EdgeSet'

Constructs a new instance from edge set fields.

Example 1:

tfgnn.EdgeSet.from_fields(
    sizes=tf.constant([3]),
    adjacency=tfgnn.Adjacency.from_indices(
        source=("paper", [1, 2, 2]), target=("paper", [0, 0, 1])))

Example 2:

 tfgnn.EdgeSet.from_fields(
     sizes=tf.constant([4]),
     adjacency=tfgnn.Adjacency.from_indices(
         source=("paper", [1, 1, 1, 2]),
         target=("author", [0, 1, 1, 3])))

Args
`features`	A mapping from feature name to feature Tensor or RaggedTensor. All feature tensors must have shape `[*graph_shape, num_edges, *feature_shape]`, where num_edge is the number of edges in the edge set (could be ragged) and feature_shape is a shape of the feature value for each edge.
`sizes`	The number of edges in each graph component. Has shape `[*graph_shape, num_components]`, where `num_components` is the number of graph components (could be ragged).
`adjacency`	One of the supported adjacency types (see adjacency.py).

Returns
An `EdgeSet` composite tensor.

get_features_dict

View source

get_features_dict() -> Dict[FieldName, Field]

Returns features copy as a dictionary.

replace_features

View source

replace_features(
    features: Mapping[FieldName, Field]
) -> '_NodeOrEdgeSet'

Returns a new instance with a new set of features.

set_shape

View source

set_shape(
    new_shape: ShapeLike
) -> 'GraphPieceSpecBase'

Enforce the common prefix shape on all the contained features.

__getitem__

View source

__getitem__(
    feature_name: FieldName
) -> Field

Indexing operator [] to access feature values by their name.