Module: tfgnn

December 14, 2023 ยท View on GitHub

View source on GitHub

Public interface for TensorFlow GNN package.

All the public symbols, data types and functions are provided from this top-level package. To use the library, you should use a single import statement, like this:

import tensorflow_gnn as tfgnn

Modules

experimental module: Experimental (unstable) parts of the public interface of TensorFlow GNN.

keras module: The tfgnn.keras package.

proto module: The protocol message (protobuf) types defined by TensorFlow GNN.

sampler module: Public interface for GNN Sampler.

Classes

class Adjacency: Stores how edges connect pairs of nodes from source and target node sets.

class AdjacencySpec: A type spec for tfgnn.Adjacency.

class Context: A composite tensor for graph context features.

class ContextSpec: A type spec for tfgnn.Context.

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

class EdgeSetSpec: A type spec for tfgnn.EdgeSet.

class Feature: The schema entry for a single feature.

class FeatureDefaultValues: Default values for graph context, node sets and edge sets features.

class GraphSchema: The top-level container for the schema of a graph dataset.

class GraphTensor: A composite tensor for heterogeneous directed graphs with features.

class GraphTensorSpec: A type spec for tfgnn.GraphTensor.

class HyperAdjacency: Stores how (hyper-)edges connect tuples of nodes from incident node sets.

class HyperAdjacencySpec: A type spec for tfgnn.HyperAdjacency.

class NodeSet: A composite tensor for node set features plus size information.

class NodeSetSpec: A type spec for tfgnn.NodeSet.

class SizeConstraints: Constraints on the number of entities in the graph.

class ValidationError: A schema validation error.

Functions

add_readout_from_first_node(...): Adds a readout structure equivalent to tfgnn.gather_first_node().

add_self_loops(...): Adds self-loops for edge_set_name EVEN if they already exist.

assert_constraints(...): Validate the shape constaints of a graph's features at runtime.

assert_satisfies_size_constraints(...): Raises InvalidArgumentError if graph_tensor exceeds size_constraints.

assert_satisfies_total_sizes(...): Raises InvalidArgumentError if graph_tensor exceeds size_constraints.

broadcast(...): Broadcasts values from nodes to edges, or from context to nodes or edges.

broadcast_context_to_edges(...): Broadcasts a context value to the edge_set edges.

broadcast_context_to_nodes(...): Broadcasts a context value to the node_set nodes.

broadcast_node_to_edges(...): Broadcasts values from nodes to incident edges.

check_compatible_with_schema_pb(...): Checks that the given spec or value is compatible with the graph schema.

check_homogeneous_graph_tensor(...): Raises ValueError when tfgnn.get_homogeneous_node_and_edge_set_name() does.

check_required_features(...): Checks the requirements of a given schema against another.

check_scalar_graph_tensor(...): Checks that graph tensor is scalar (has rank 0).

combine_values(...): Combines a list of tensors into one (by concatenation or otherwise).

convert_to_line_graph(...): Obtain a graph's line graph.

create_graph_spec_from_schema_pb(...): Converts a graph schema proto message to a scalar GraphTensorSpec.

create_schema_pb_from_graph_spec(...): Converts scalar GraphTensorSpec to a graph schema proto message.

dataset_filter_with_summary(...): Dataset filter with a summary for the fraction of dataset elements removed.

dataset_from_generator(...): Creates dataset from generator of any nest of scalar graph pieces.

disable_graph_tensor_validation(...): Disables both static and runtime checks of graph tensors.

disable_graph_tensor_validation_at_runtime(...): Disables runtime checks (tf.debugging.Assert) of graph tensors.

enable_graph_tensor_validation(...): Enables static checks of graph tensors.

enable_graph_tensor_validation_at_runtime(...): Enables both static and runtime checks of graph tensors.

find_tight_size_constraints(...): Returns smallest possible size constraints that allow dataset padding.

gather_first_node(...): Gathers feature value from the first node of each graph component.

get_aux_type_prefix(...): Returns type prefix of aux node or edge set names, or None if non-aux.

get_homogeneous_node_and_edge_set_name(...): Returns the sole node_set_name, edge_set_name or raises ValueError.

get_io_spec(...): Returns tf.io parsing features for GraphTensorSpec type spec.

get_registered_reduce_operation_names(...): Returns the registered list of supported reduce operation names.

graph_tensor_to_values(...): Convert an eager GraphTensor to a mapping of mappings of PODTs.

homogeneous(...): Constructs a homogeneous GraphTensor with node features and one edge_set.

is_dense_tensor(...): Returns whether a tensor (TF or Keras) is a Tensor.

is_graph_tensor(...): Returns whether value is a GraphTensor (possibly wrapped for Keras).

is_ragged_tensor(...): Returns whether a tensor (TF or Keras) is a RaggedTensor.

iter_features(...): Utility function to iterate over the features of a graph schema.

iter_sets(...): Utility function to iterate over all the sets present in a graph schema.

learn_fit_or_skip_size_constraints(...): Learns the optimal size constraints for the fixed size batching with retry.

mask_edges(...): Creates a GraphTensor after applying edge_mask over the specified edge-set.

node_degree(...): Returns the degree of each node w.r.t. one side of an edge set.

pad_to_total_sizes(...): Pads graph tensor to the total sizes by inserting fake graph components.

parse_example(...): Parses a batch of serialized Example protos into a single GraphTensor.

parse_schema(...): Parse a schema from text-formatted protos.

parse_single_example(...): Parses a single serialized Example proto into a single GraphTensor.

pool(...): Pools values from edges to nodes, or from nodes or edges to context.

pool_edges_to_context(...): Aggregates (pools) edge values to graph context.

pool_edges_to_node(...): Aggregates (pools) edge values to incident nodes.

pool_neighbors_to_node(...): Aggregates (pools) neighbor node values along one or more edge sets.

pool_neighbors_to_node_feature(...): Aggregates (pools) sender node feature to receiver nodes feature.

pool_nodes_to_context(...): Aggregates (pools) node values to graph context.

random_graph_tensor(...): Generate a graph tensor from a spec, with random features.

read_schema(...): Read a proto schema from a file with text-formatted contents.

reorder_nodes(...): Reorders nodes within node sets according to indices.

reverse_tag(...): Flips tfgnn.SOURCE to tfgnn.TARGET and vice versa.

satisfies_size_constraints(...): Returns whether the input graph_tensor satisfies total_sizes.

satisfies_total_sizes(...): Returns whether the input graph_tensor satisfies total_sizes.

shuffle_features_globally(...): Shuffles context, node set and edge set features of a scalar GraphTensor.

shuffle_nodes(...): Randomly reorders nodes of given node sets, within each graph component.

softmax(...): Computes softmax over a many-to-one relationship in a GraphTensor.

softmax_edges_per_node(...): Returns softmax() of edge values per common node_tag node.

structured_readout(...): Reads out a feature value from select nodes (or edges) in a graph.

structured_readout_into_feature(...): Reads out a feature value from select nodes (or edges) in a graph.

validate_graph_tensor_for_readout(...): Checks graph supports structured_readout() from required_keys.

validate_graph_tensor_spec_for_readout(...): Checks graph_spec supports structured_readout() from required_keys.

validate_schema(...): Validates the correctness of a graph schema instance.

write_example(...): Encode an eager GraphTensor to a tf.train.Example proto.

write_schema(...): Write a GraphSchema to a text-formatted proto file.

Type Aliases

Field

FieldOrFields

FieldSpec

Fields

FieldsSpec

IncidentNodeOrContextTag

CONTEXT 'context'
EDGES 'edges'
HIDDEN_STATE 'hidden_state'
NODES 'nodes'
SIZE_NAME '#size'
SOURCE 0
SOURCE_NAME '#source'
TARGET 1
TARGET_NAME '#target'
**version** '1.0.0.dev3'