tfgnn.EdgeSet

View source on GitHub

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

tfgnn.EdgeSet(
    data: Data, spec: 'GraphPieceSpecBase'
)

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.

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 dtype for graph items indexing. One of tf.int32 or tf.int64.
num_components	The number of graph components for each graph.
rank	The rank of this Tensor. Guaranteed not to be None.
row_splits_dtype	The dtype for ragged row partitions. One of tf.int32 or tf.int64.
shape	A possibly-partial shape specification for this Tensor. The returned tf.TensorShape is guaranteed to have a known rank and no unknown dimensions except possibly the outermost.
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,
    validate: Optional[bool] = None
) -> '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).
validate	If true, use tf.assert ops to inspect the shapes of each field and check at runtime that they form a valid EdgeSet. The default behavior is set by the disable_graph_tensor_validation_at_runtime() and enable_graph_tensor_validation_at_runtime().

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
) -> 'GraphPieceBase'

Deprecated. Use with_shape().

with_indices_dtype

View source

with_indices_dtype(
    dtype: tf.dtypes.DType
) -> 'GraphPieceBase'

Returns a copy of this piece with the given indices dtype.

with_row_splits_dtype

View source

with_row_splits_dtype(
    dtype: tf.dtypes.DType
) -> 'GraphPieceBase'

Returns a copy of this piece with the given row splits dtype.

with_shape

View source

with_shape(
    new_shape: ShapeLike
) -> 'GraphPieceBase'

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.