From dd001a484975bd2548441733cc0ce5805c12ffe3 Mon Sep 17 00:00:00 2001 From: Oleksandr Ferludin Date: Wed, 6 Dec 2023 10:33:06 -0800 Subject: [PATCH] Proposal for updated API Docs PiperOrigin-RevId: 588473240 --- .../python/models/contrastive_losses.md | 69 + .../contrastive_losses/AllSvdMetrics.md | 202 +++ .../contrastive_losses/BarlowTwinsTask.md | 168 +++ .../contrastive_losses/ContrastiveLossTask.md | 186 +++ .../contrastive_losses/CorruptionSpec.md | 87 ++ .../models/contrastive_losses/Corruptor.md | 58 + .../DeepGraphInfomaxLogits.md | 17 + .../DeepGraphInfomaxTask.md | 165 +++ .../contrastive_losses/DropoutFeatures.md | 56 + .../ShuffleFeaturesGlobally.md | 60 + .../TripletEmbeddingSquaredDistances.md | 17 + .../contrastive_losses/TripletLossTask.md | 200 +++ .../models/contrastive_losses/VicRegTask.md | 169 +++ .../models/contrastive_losses/all_symbols.md | 24 + .../models/contrastive_losses/coherence.md | 69 + .../contrastive_losses/numerical_rank.md | 70 + .../pseudo_condition_number.md | 69 + .../models/contrastive_losses/rankme.md | 77 ++ .../contrastive_losses/self_clustering.md | 63 + .../docs/api_docs/python/models/gat_v2.md | 13 +- .../python/models/gat_v2/GATv2Conv.md | 160 +-- .../python/models/gat_v2/GATv2EdgePool.md | 34 +- .../models/gat_v2/GATv2HomGraphUpdate.md | 36 +- .../models/gat_v2/GATv2MPNNGraphUpdate.md | 65 +- .../gat_v2/graph_update_from_config_dict.md | 45 +- .../gat_v2/graph_update_get_config_dict.md | 13 +- .../docs/api_docs/python/models/gcn.md | 13 +- .../api_docs/python/models/gcn/GCNConv.md | 60 +- .../python/models/gcn/GCNHomGraphUpdate.md | 36 +- .../docs/api_docs/python/models/graph_sage.md | 13 +- .../graph_sage/GCNGraphSAGENodeSetUpdate.md | 47 +- .../graph_sage/GraphSAGEAggregatorConv.md | 52 +- .../models/graph_sage/GraphSAGEGraphUpdate.md | 54 +- .../models/graph_sage/GraphSAGENextState.md | 36 +- .../models/graph_sage/GraphSAGEPoolingConv.md | 60 +- .../docs/api_docs/python/models/mt_albis.md | 13 +- .../models/mt_albis/MtAlbisGraphUpdate.md | 108 +- .../mt_albis/graph_update_from_config_dict.md | 42 +- .../mt_albis/graph_update_get_config_dict.md | 13 +- .../python/models/multi_head_attention.md | 13 +- .../MultiHeadAttentionConv.md | 266 ++-- .../MultiHeadAttentionEdgePool.md | 36 +- .../MultiHeadAttentionHomGraphUpdate.md | 34 +- .../MultiHeadAttentionMPNNGraphUpdate.md | 65 +- .../graph_update_from_config_dict.md | 49 +- .../graph_update_get_config_dict.md | 13 +- .../api_docs/python/models/vanilla_mpnn.md | 13 +- .../vanilla_mpnn/VanillaMPNNGraphUpdate.md | 57 +- .../graph_update_from_config_dict.md | 47 +- .../graph_update_get_config_dict.md | 13 +- tensorflow_gnn/docs/api_docs/python/runner.md | 173 +++ .../api_docs/python/runner/ContextLabelFn.md | 17 + .../api_docs/python/runner/DatasetProvider.md | 27 + .../python/runner/DotProductLinkPrediction.md | 200 +++ .../python/runner/FitOrSkipPadding.md | 49 + .../runner/GraphBinaryClassification.md | 222 +++ .../python/runner/GraphMeanAbsoluteError.md | 195 +++ .../GraphMeanAbsolutePercentageError.md | 195 +++ .../python/runner/GraphMeanSquaredError.md | 195 +++ .../runner/GraphMeanSquaredLogScaledError.md | 155 +++ .../GraphMeanSquaredLogarithmicError.md | 195 +++ .../runner/GraphMulticlassClassification.md | 239 ++++ .../python/runner/GraphTensorPadding.md | 37 + .../python/runner/GraphTensorProcessorFn.md | 26 + .../runner/HadamardProductLinkPrediction.md | 203 +++ .../runner/IntegratedGradientsExporter.md | 127 ++ .../python/runner/KerasModelExporter.md | 121 ++ .../api_docs/python/runner/KerasTrainer.md | 263 ++++ .../runner/KerasTrainerCheckpointOptions.md | 108 ++ .../python/runner/KerasTrainerOptions.md | 90 ++ .../docs/api_docs/python/runner/Loss.md | 17 + .../docs/api_docs/python/runner/Losses.md | 16 + .../docs/api_docs/python/runner/Metrics.md | 17 + .../api_docs/python/runner/ModelExporter.md | 53 + .../python/runner/NodeBinaryClassification.md | 237 ++++ .../runner/NodeMulticlassClassification.md | 254 ++++ .../python/runner/ParameterServerStrategy.md | 1082 +++++++++++++++ .../python/runner/PassthruDatasetProvider.md | 40 + .../runner/PassthruSampleDatasetsProvider.md | 46 + .../api_docs/python/runner/Predictions.md | 17 + .../runner/RootNodeBinaryClassification.md | 216 +++ .../api_docs/python/runner/RootNodeLabelFn.md | 20 + .../runner/RootNodeMeanAbsoluteError.md | 194 +++ .../RootNodeMeanAbsoluteLogarithmicError.md | 152 ++ .../RootNodeMeanAbsolutePercentageError.md | 194 +++ .../python/runner/RootNodeMeanSquaredError.md | 194 +++ .../RootNodeMeanSquaredLogScaledError.md | 155 +++ .../RootNodeMeanSquaredLogarithmicError.md | 194 +++ .../RootNodeMulticlassClassification.md | 233 ++++ .../docs/api_docs/python/runner/RunResult.md | 71 + .../runner/SampleTFRecordDatasetsProvider.md | 229 +++ .../python/runner/SimpleDatasetProvider.md | 97 ++ .../runner/SimpleSampleDatasetsProvider.md | 238 ++++ .../python/runner/SubmoduleExporter.md | 125 ++ .../python/runner/TFDataServiceConfig.md | 66 + .../python/runner/TFRecordDatasetProvider.md | 93 ++ .../api_docs/python/runner/TPUStrategy.md | 1228 +++++++++++++++++ .../docs/api_docs/python/runner/Task.md | 179 +++ .../api_docs/python/runner/TightPadding.md | 47 + .../docs/api_docs/python/runner/Trainer.md | 100 ++ .../api_docs/python/runner/all_symbols.md | 62 + .../api_docs/python/runner/export_model.md | 77 ++ .../python/runner/incrementing_model_dir.md | 52 + .../python/runner/integrated_gradients.md | 95 ++ .../python/runner/one_node_per_component.md | 17 + .../docs/api_docs/python/runner/run.md | 275 ++++ tensorflow_gnn/docs/api_docs/python/tfgnn.md | 234 ++-- .../docs/api_docs/python/tfgnn/Adjacency.md | 174 +-- .../api_docs/python/tfgnn/AdjacencySpec.md | 218 +-- .../docs/api_docs/python/tfgnn/Context.md | 192 +-- .../docs/api_docs/python/tfgnn/ContextSpec.md | 189 +-- .../docs/api_docs/python/tfgnn/EdgeSet.md | 175 ++- .../docs/api_docs/python/tfgnn/EdgeSetSpec.md | 197 +-- .../docs/api_docs/python/tfgnn/Feature.md | 75 - .../python/tfgnn/FeatureDefaultValues.md | 32 +- .../docs/api_docs/python/tfgnn/Field.md | 10 +- .../api_docs/python/tfgnn/FieldOrFields.md | 10 +- .../docs/api_docs/python/tfgnn/FieldSpec.md | 10 +- .../docs/api_docs/python/tfgnn/Fields.md | 10 +- .../docs/api_docs/python/tfgnn/FieldsSpec.md | 10 +- .../docs/api_docs/python/tfgnn/GraphSchema.md | 90 -- .../docs/api_docs/python/tfgnn/GraphTensor.md | 289 ++-- .../api_docs/python/tfgnn/GraphTensorSpec.md | 192 +-- .../api_docs/python/tfgnn/HyperAdjacency.md | 168 +-- .../python/tfgnn/HyperAdjacencySpec.md | 200 +-- .../python/tfgnn/IncidentNodeOrContextTag.md | 10 +- .../docs/api_docs/python/tfgnn/NodeSet.md | 170 ++- .../docs/api_docs/python/tfgnn/NodeSetSpec.md | 193 +-- .../api_docs/python/tfgnn/SizeConstraints.md | 36 +- .../api_docs/python/tfgnn/ValidationError.md | 16 +- .../tfgnn/add_readout_from_first_node.md | 23 +- .../api_docs/python/tfgnn/add_self_loops.md | 49 +- .../docs/api_docs/python/tfgnn/all_symbols.md | 21 +- .../python/tfgnn/assert_constraints.md | 26 +- .../assert_satisfies_size_constraints.md | 38 +- .../docs/api_docs/python/tfgnn/broadcast.md | 34 +- .../tfgnn/broadcast_context_to_edges.md | 34 +- .../tfgnn/broadcast_context_to_nodes.md | 34 +- .../python/tfgnn/broadcast_node_to_edges.md | 44 +- .../tfgnn/check_compatible_with_schema_pb.md | 25 +- .../tfgnn/check_homogeneous_graph_tensor.md | 14 +- .../python/tfgnn/check_required_features.md | 48 +- .../python/tfgnn/check_scalar_graph_tensor.md | 15 +- .../api_docs/python/tfgnn/combine_values.md | 41 +- .../python/tfgnn/convert_to_line_graph.md | 31 +- .../tfgnn/create_graph_spec_from_schema_pb.md | 39 +- .../tfgnn/create_schema_pb_from_graph_spec.md | 36 +- .../tfgnn/dataset_filter_with_summary.md | 36 +- .../python/tfgnn/dataset_from_generator.md | 22 +- .../tfgnn/disable_graph_tensor_validation.md | 19 + ...able_graph_tensor_validation_at_runtime.md | 15 + .../tfgnn/enable_graph_tensor_validation.md | 15 + ...able_graph_tensor_validation_at_runtime.md | 15 + .../api_docs/python/tfgnn/experimental.md | 13 +- .../tfgnn/find_tight_size_constraints.md | 42 +- .../python/tfgnn/gather_first_node.md | 44 +- .../python/tfgnn/get_aux_type_prefix.md | 19 +- .../get_homogeneous_node_and_edge_set_name.md | 25 +- .../docs/api_docs/python/tfgnn/get_io_spec.md | 35 +- .../get_registered_reduce_operation_names.md | 13 +- .../python/tfgnn/graph_tensor_to_values.md | 24 +- .../docs/api_docs/python/tfgnn/homogeneous.md | 36 +- .../api_docs/python/tfgnn/is_dense_tensor.md | 13 +- .../api_docs/python/tfgnn/is_graph_tensor.md | 15 +- .../api_docs/python/tfgnn/is_ragged_tensor.md | 13 +- .../api_docs/python/tfgnn/iter_features.md | 36 +- .../docs/api_docs/python/tfgnn/iter_sets.md | 36 +- .../docs/api_docs/python/tfgnn/keras.md | 20 +- .../python/tfgnn/keras/ConvGNNBuilder.md | 56 +- .../python/tfgnn/keras/clone_initializer.md | 23 +- .../api_docs/python/tfgnn/keras/layers.md | 59 +- .../keras/layers/AddReadoutFromFirstNode.md | 25 +- .../python/tfgnn/keras/layers/AddSelfLoops.md | 13 +- .../keras/layers/AnyToAnyConvolutionBase.md | 228 ++- .../python/tfgnn/keras/layers/Broadcast.md | 46 +- .../tfgnn/keras/layers/ContextUpdate.md | 59 +- .../tfgnn/keras/layers/EdgeSetUpdate.md | 40 +- .../python/tfgnn/keras/layers/GraphUpdate.md | 44 +- .../python/tfgnn/keras/layers/ItemDropout.md | 39 +- .../tfgnn/keras/layers/MakeEmptyFeature.md | 28 +- .../python/tfgnn/keras/layers/MapFeatures.md | 106 +- .../tfgnn/keras/layers/NextStateFromConcat.md | 22 +- .../tfgnn/keras/layers/NodeSetUpdate.md | 54 +- .../tfgnn/keras/layers/PadToTotalSizes.md | 30 +- .../python/tfgnn/keras/layers/ParseExample.md | 18 +- .../tfgnn/keras/layers/ParseSingleExample.md | 18 +- .../python/tfgnn/keras/layers/Pool.md | 59 +- .../python/tfgnn/keras/layers/Readout.md | 64 +- .../tfgnn/keras/layers/ReadoutFirstNode.md | 54 +- .../keras/layers/ReadoutNamedIntoFeature.md | 57 +- .../tfgnn/keras/layers/ResidualNextState.md | 26 +- .../python/tfgnn/keras/layers/SimpleConv.md | 148 +- .../keras/layers/SingleInputNextState.md | 17 +- .../tfgnn/keras/layers/StructuredReadout.md | 31 +- .../learn_fit_or_skip_size_constraints.md | 58 +- .../docs/api_docs/python/tfgnn/mask_edges.md | 27 +- .../docs/api_docs/python/tfgnn/node_degree.md | 30 +- .../python/tfgnn/pad_to_total_sizes.md | 66 +- .../api_docs/python/tfgnn/parse_example.md | 38 +- .../api_docs/python/tfgnn/parse_schema.md | 26 +- .../python/tfgnn/parse_single_example.md | 35 +- .../docs/api_docs/python/tfgnn/pool.md | 46 +- .../python/tfgnn/pool_edges_to_context.md | 40 +- .../python/tfgnn/pool_edges_to_node.md | 46 +- .../python/tfgnn/pool_neighbors_to_node.md | 110 ++ .../tfgnn/pool_neighbors_to_node_feature.md | 105 ++ .../python/tfgnn/pool_nodes_to_context.md | 40 +- .../docs/api_docs/python/tfgnn/proto.md | 68 + .../api_docs/python/tfgnn/proto/BigQuery.md | 91 ++ .../python/tfgnn/proto/BigQuery/TableSpec.md | 40 + .../api_docs/python/tfgnn/proto/Context.md | 36 + .../api_docs/python/tfgnn/proto/EdgeSet.md | 64 + .../api_docs/python/tfgnn/proto/Feature.md | 58 + .../python/tfgnn/proto/GraphSchema.md | 58 + .../api_docs/python/tfgnn/proto/Metadata.md | 54 + .../python/tfgnn/proto/Metadata/KeyValue.md | 33 + .../api_docs/python/tfgnn/proto/NodeSet.md | 50 + .../api_docs/python/tfgnn/proto/OriginInfo.md | 36 + .../python/tfgnn/random_graph_tensor.md | 57 +- .../docs/api_docs/python/tfgnn/read_schema.md | 28 +- .../api_docs/python/tfgnn/reorder_nodes.md | 47 +- .../docs/api_docs/python/tfgnn/reverse_tag.md | 15 +- .../docs/api_docs/python/tfgnn/sampler.md | 29 +- .../python/tfgnn/sampler/SamplingOp.md | 33 +- .../python/tfgnn/sampler/SamplingSpec.md | 28 +- .../tfgnn/sampler/SamplingSpecBuilder.md | 27 +- .../tfgnn/sampler/make_sampling_spec_tree.md | 35 +- .../tfgnn/satisfies_size_constraints.md | 28 +- .../python/tfgnn/shuffle_features_globally.md | 29 +- .../api_docs/python/tfgnn/shuffle_nodes.md | 30 +- .../docs/api_docs/python/tfgnn/softmax.md | 41 +- .../python/tfgnn/softmax_edges_per_node.md | 13 +- .../python/tfgnn/structured_readout.md | 33 +- .../tfgnn/structured_readout_into_feature.md | 51 +- .../validate_graph_tensor_for_readout.md | 29 +- .../validate_graph_tensor_spec_for_readout.md | 25 +- .../api_docs/python/tfgnn/validate_schema.md | 34 +- .../api_docs/python/tfgnn/write_example.md | 36 +- .../api_docs/python/tfgnn/write_schema.md | 28 +- .../models/multi_head_attention/layers.py | 28 +- 240 files changed, 15611 insertions(+), 4234 deletions(-) create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/AllSvdMetrics.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/BarlowTwinsTask.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ContrastiveLossTask.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/CorruptionSpec.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/Corruptor.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxLogits.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxTask.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DropoutFeatures.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ShuffleFeaturesGlobally.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletEmbeddingSquaredDistances.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletLossTask.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/VicRegTask.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/all_symbols.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/coherence.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/numerical_rank.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/pseudo_condition_number.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/rankme.md create mode 100644 tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/self_clustering.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/ContextLabelFn.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/DatasetProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/DotProductLinkPrediction.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/FitOrSkipPadding.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphBinaryClassification.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsoluteError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsolutePercentageError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogScaledError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogarithmicError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphMulticlassClassification.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphTensorPadding.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/GraphTensorProcessorFn.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/HadamardProductLinkPrediction.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/IntegratedGradientsExporter.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/KerasModelExporter.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/KerasTrainer.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerCheckpointOptions.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerOptions.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/Loss.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/Losses.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/Metrics.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/ModelExporter.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/NodeBinaryClassification.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/NodeMulticlassClassification.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/ParameterServerStrategy.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/PassthruDatasetProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/PassthruSampleDatasetsProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/Predictions.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeBinaryClassification.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeLabelFn.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteLogarithmicError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsolutePercentageError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogScaledError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogarithmicError.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RootNodeMulticlassClassification.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/RunResult.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/SampleTFRecordDatasetsProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/SimpleDatasetProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/SimpleSampleDatasetsProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/SubmoduleExporter.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/TFDataServiceConfig.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/TFRecordDatasetProvider.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/TPUStrategy.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/Task.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/TightPadding.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/Trainer.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/all_symbols.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/export_model.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/incrementing_model_dir.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/integrated_gradients.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/one_node_per_component.md create mode 100644 tensorflow_gnn/docs/api_docs/python/runner/run.md delete mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/Feature.md delete mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/GraphSchema.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation_at_runtime.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation_at_runtime.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node_feature.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery/TableSpec.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Context.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/EdgeSet.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Feature.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/GraphSchema.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata/KeyValue.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/NodeSet.md create mode 100644 tensorflow_gnn/docs/api_docs/python/tfgnn/proto/OriginInfo.md diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses.md new file mode 100644 index 00000000..0eb3460d --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses.md @@ -0,0 +1,69 @@ +# Module: contrastive_losses + + + + + View source +on GitHub + +Contrastive losses. + +Users of TF-GNN can use these layers by importing them next to the core library: + +```python +import tensorflow_gnn as tfgnn +from tensorflow_gnn.models import contrastive_losses +``` + +## Classes + +[`class AllSvdMetrics`](./contrastive_losses/AllSvdMetrics.md): Computes +multiple metrics for representations using one SVD call. + +[`class BarlowTwinsTask`](./contrastive_losses/BarlowTwinsTask.md): A Barlow +Twins (BT) Task. + +[`class ContrastiveLossTask`](./contrastive_losses/ContrastiveLossTask.md): Base +class for unsupervised contrastive representation learning tasks. + +[`class CorruptionSpec`](./contrastive_losses/CorruptionSpec.md): Class for +defining corruption specification. + +[`class Corruptor`](./contrastive_losses/Corruptor.md): Base class for graph +corruptor. + +[`class DeepGraphInfomaxLogits`](./contrastive_losses/DeepGraphInfomaxLogits.md): +Computes clean and corrupted logits for Deep Graph Infomax (DGI). + +[`class DeepGraphInfomaxTask`](./contrastive_losses/DeepGraphInfomaxTask.md): A +Deep Graph Infomax (DGI) Task. + +[`class DropoutFeatures`](./contrastive_losses/DropoutFeatures.md): Base class +for graph corruptor. + +[`class ShuffleFeaturesGlobally`](./contrastive_losses/ShuffleFeaturesGlobally.md): +A corruptor that shuffles features. + +[`class TripletEmbeddingSquaredDistances`](./contrastive_losses/TripletEmbeddingSquaredDistances.md): +Computes embeddings distance between positive and negative pairs. + +[`class TripletLossTask`](./contrastive_losses/TripletLossTask.md): The triplet +loss task. + +[`class VicRegTask`](./contrastive_losses/VicRegTask.md): A VICReg Task. + +## Functions + +[`coherence(...)`](./contrastive_losses/coherence.md): Coherence metric +implementation. + +[`numerical_rank(...)`](./contrastive_losses/numerical_rank.md): Numerical rank +implementation. + +[`pseudo_condition_number(...)`](./contrastive_losses/pseudo_condition_number.md): +Pseudo-condition number metric implementation. + +[`rankme(...)`](./contrastive_losses/rankme.md): RankMe metric implementation. + +[`self_clustering(...)`](./contrastive_losses/self_clustering.md): +Self-clustering metric implementation. diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/AllSvdMetrics.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/AllSvdMetrics.md new file mode 100644 index 00000000..5d7964d3 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/AllSvdMetrics.md @@ -0,0 +1,202 @@ +# contrastive_losses.AllSvdMetrics + + + + + View source +on GitHub + +Computes multiple metrics for representations using one SVD call. + +
+contrastive_losses.AllSvdMetrics(
+    *args, **kwargs
+)
+
+ + + +Refer to https://arxiv.org/abs/2305.16562 for more details. + + + + + + + + + + + + + + + + + +

Args

+fns + +a mapping from a metric name to a Callable that accepts +representations as well as the result of their SVD decomposition. +Currently only singular values are passed. +
+y_pred_transform_fn + +a function to extract clean representations +from model predictions. By default, no transformation is applied. +
+name + +Name for the metric class, used for Keras bookkeeping. +
+ +## Methods + +

merge_state

+ +
+merge_state(
+    metrics
+)
+
+ +Merges the state from one or more metrics. + +This method can be used by distributed systems to merge the state computed by +different metric instances. Typically the state will be stored in the form of +the metric's weights. For example, a tf.keras.metrics.Mean metric contains a +list of two weight values: a total and a count. If there were two instances of a +tf.keras.metrics.Accuracy that each independently aggregated partial state for +an overall accuracy calculation, these two metric's states could be combined as +follows: + +``` +>>> m1 = tf.keras.metrics.Accuracy() +>>> _ = m1.update_state([[1], [2]], [[0], [2]]) +``` + +``` +>>> m2 = tf.keras.metrics.Accuracy() +>>> _ = m2.update_state([[3], [4]], [[3], [4]]) +``` + +``` +>>> m2.merge_state([m1]) +>>> m2.result().numpy() +0.75 +``` + + + + + + + + + + + +
Args
+metrics + +an iterable of metrics. The metrics must have compatible +state. +
+ + + + + + + + + + + +
Raises
+ValueError + +If the provided iterable does not contain metrics matching +the metric's required specifications. +
+ +

reset_state

+ +View +source + +
+reset_state() -> None
+
+ +Resets all of the metric state variables. + +This function is called between epochs/steps, when a metric is evaluated during +training. + +

result

+ +View +source + +
+result() -> Mapping[str, tf.Tensor]
+
+ +Computes and returns the scalar metric value tensor or a dict of scalars. + +Result computation is an idempotent operation that simply calculates the metric +value using the state variables. + + + + + + + + + + +
Returns
+A scalar tensor, or a dictionary of scalar tensors. +
+ +

update_state

+ +View +source + +
+update_state(
+    _, y_pred: tf.Tensor, sample_weight=None
+) -> None
+
+ +Accumulates statistics for the metric. + +Note: This function is executed as a graph function in graph mode. This means: +a) Operations on the same resource are executed in textual order. This should +make it easier to do things like add the updated value of a variable to another, +for example. b) You don't need to worry about collecting the update ops to +execute. All update ops added to the graph by this function will be executed. As +a result, code should generally work the same way with graph or eager execution. + + + + + + + + + + + + +
Args
*args + +
+**kwargs + +A mini-batch of inputs to the Metric. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/BarlowTwinsTask.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/BarlowTwinsTask.md new file mode 100644 index 00000000..8ef94ef3 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/BarlowTwinsTask.md @@ -0,0 +1,168 @@ +# contrastive_losses.BarlowTwinsTask + + + + + View source +on GitHub + +A Barlow Twins (BT) Task. + +Inherits From: +[`ContrastiveLossTask`](../contrastive_losses/ContrastiveLossTask.md) + +
+contrastive_losses.BarlowTwinsTask(
+    *args,
+    lambda_: Optional[Union[tf.Tensor, float]] = None,
+    normalize_batch: bool = True,
+    **kwargs
+)
+
+ + + + + + + + + + + + + + + + + + + + + + + + + + + +

Args

+node_set_name + +Name of the node set for readout. +
+feature_name + +Feature name for readout. +
+representations_layer_name + +Layer name for uncorrupted representations. +
+corruptor + +Corruptor instance for creating negative samples. If not +specified, we use ShuffleFeaturesGlobally by default. +
+projector_units + +Sequence of layer sizes for projector network. +Projectors prevent dimensional collapse, but can hinder training for +easy corruptions. For more details, see +https://arxiv.org/abs/2304.12210. +
+seed + +Random seed for the default corruptor (ShuffleFeaturesGlobally). +
+ +## Methods + +

losses

+ +View +source + +
+losses() -> runner.Losses
+
+ +Returns arbitrary task specific losses. + +

make_contrastive_layer

+ +View +source + +
+make_contrastive_layer() -> tf.keras.layers.Layer
+
+ +Returns the layer contrasting clean outputs with the correupted ones. + +

metrics

+ +View +source + +
+metrics() -> runner.Metrics
+
+ +Returns arbitrary task specific metrics. + +

predict

+ +View +source + +
+predict(
+    *args
+) -> runner.Predictions
+
+ +Apply a readout head for use with various contrastive losses. + + + + + + + + + + + +
Args
+*args + +A tuple of (clean, corrupted) tfgnn.GraphTensors. +
+ + + + + + + + + + +
Returns
+The logits for some contrastive loss as produced by the implementing +subclass. +
+ +

preprocess

+ +View +source + +
+preprocess(
+    inputs: GraphTensor
+) -> tuple[Sequence[GraphTensor], runner.Predictions]
+
+ +Applies a `Corruptor` and returns empty pseudo-labels. diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ContrastiveLossTask.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ContrastiveLossTask.md new file mode 100644 index 00000000..f25e9ec3 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ContrastiveLossTask.md @@ -0,0 +1,186 @@ +# contrastive_losses.ContrastiveLossTask + + + + + View source +on GitHub + +Base class for unsupervised contrastive representation learning tasks. + +
+contrastive_losses.ContrastiveLossTask(
+    node_set_name: str,
+    *,
+    feature_name: str = tfgnn.HIDDEN_STATE,
+    representations_layer_name: Optional[str] = None,
+    corruptor: Optional[layers.Corruptor] = None,
+    projector_units: Optional[Sequence[int]] = None,
+    seed: Optional[int] = None
+)
+
+ + + +The process is separated into preprocessing and contrastive parts, with the +focus on reusability of individual components. The `preprocess` produces input +GraphTensors to be used with the `predict` as well as labels for the task. The +default `predict` method implementation expects a pair of positive and negative +GraphTensors. There are multiple ways proposed in the literature to learn +representations based on the activations - we achieve that by using custom +losses. + +Any subclass must implement `make_contrastive_layer` method, which produces the +final prediction outputs. + +If the loss involves labels for each example, subclasses should leverage +`losses` and `metrics` methods to specify task's losses. When the loss only +involves model outputs, `make_contrastive_layer` should output both positive and +perturb examples, and the `losses` should use pseudolabels. + +Any model-specific preprocessing should be implemented in the `preprocess`. + + + + + + + + + + + + + + + + + + + + + + + + + + +

Args

+node_set_name + +Name of the node set for readout. +
+feature_name + +Feature name for readout. +
+representations_layer_name + +Layer name for uncorrupted representations. +
+corruptor + +Corruptor instance for creating negative samples. If not +specified, we use ShuffleFeaturesGlobally by default. +
+projector_units + +Sequence of layer sizes for projector network. +Projectors prevent dimensional collapse, but can hinder training for +easy corruptions. For more details, see +https://arxiv.org/abs/2304.12210. +
+seed + +Random seed for the default corruptor (ShuffleFeaturesGlobally). +
+ +## Methods + +

losses

+ +
+@abc.abstractmethod
+losses() -> Losses
+
+ +Returns arbitrary task specific losses. + +

make_contrastive_layer

+ +View +source + +
+@abc.abstractmethod
+make_contrastive_layer() -> tf.keras.layers.Layer
+
+ +Returns the layer contrasting clean outputs with the correupted ones. + +

metrics

+ +View +source + +
+metrics() -> runner.Metrics
+
+ +Returns arbitrary task specific metrics. + +

predict

+ +View +source + +
+predict(
+    *args
+) -> runner.Predictions
+
+ +Apply a readout head for use with various contrastive losses. + + + + + + + + + + + +
Args
+*args + +A tuple of (clean, corrupted) tfgnn.GraphTensors. +
+ + + + + + + + + + +
Returns
+The logits for some contrastive loss as produced by the implementing +subclass. +
+ +

preprocess

+ +View +source + +
+preprocess(
+    inputs: GraphTensor
+) -> tuple[Sequence[GraphTensor], runner.Predictions]
+
+ +Applies a `Corruptor` and returns empty pseudo-labels. diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/CorruptionSpec.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/CorruptionSpec.md new file mode 100644 index 00000000..ead30912 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/CorruptionSpec.md @@ -0,0 +1,87 @@ +# contrastive_losses.CorruptionSpec + + + + + View source +on GitHub + +Class for defining corruption specification. + +
+contrastive_losses.CorruptionSpec(
+    node_set_corruption: NodeCorruptionSpec = dataclasses.field(default_factory=dict),
+    edge_set_corruption: EdgeCorruptionSpec = dataclasses.field(default_factory=dict),
+    context_corruption: ContextCorruptionSpec = dataclasses.field(default_factory=dict)
+)
+
+ + + +This has three fields for specifying the corruption behavior of node-, edge-, +and context sets. + +A value of the key "*" is a wildcard value that is used for either all features +or all node/edge sets. + +#### Some example usages: + +Want: corrupt everything with parameter 1.0. Solution: either set default to 1.0 +or set all corruption specs to `{"*": 1.}`. + +Want: corrupt all context features with parameter 1.0 except for "feat", which +should not be corrupted. Solution: set `context_corruption` to `{"feat": 0., +"*": 1.}` + + + + + + + + + + + + + + + + + +

Attributes

+node_set_corruption + +Dataclass field +
+edge_set_corruption + +Dataclass field +
+context_corruption + +Dataclass field +
+ +## Methods + +

with_default

+ +View +source + +
+with_default(
+    default: T
+)
+
+ +

__eq__

+ +
+__eq__(
+    other
+)
+
+ +Return self==value. diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/Corruptor.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/Corruptor.md new file mode 100644 index 00000000..89a71bb3 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/Corruptor.md @@ -0,0 +1,58 @@ +# contrastive_losses.Corruptor + + + + + View source +on GitHub + +Base class for graph corruptor. + +
+contrastive_losses.Corruptor(
+    corruption_spec: Optional[CorruptionSpec[T]] = None,
+    *,
+    corruption_fn: Callable[[tfgnn.Field, T], tfgnn.Field],
+    default: Optional[T] = None,
+    **kwargs
+)
+
+ + + + + + + + + + + + + + + + + + + + + +

Args

+corruption_spec + +A spec for corruption application. +
+corruption_fn + +Corruption function. +
+default + +Global application default of the corruptor. This is only used +when corruption_spec is None. +
+**kwargs + +Additional keyword arguments. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxLogits.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxLogits.md new file mode 100644 index 00000000..60a512a3 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxLogits.md @@ -0,0 +1,17 @@ +# contrastive_losses.DeepGraphInfomaxLogits + + + + + View source +on GitHub + +Computes clean and corrupted logits for Deep Graph Infomax (DGI). + +
+contrastive_losses.DeepGraphInfomaxLogits(
+    trainable=True, name=None, dtype=None, dynamic=False, **kwargs
+)
+
+ + diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxTask.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxTask.md new file mode 100644 index 00000000..1b35ea21 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DeepGraphInfomaxTask.md @@ -0,0 +1,165 @@ +# contrastive_losses.DeepGraphInfomaxTask + + + + + View source +on GitHub + +A Deep Graph Infomax (DGI) Task. + +Inherits From: +[`ContrastiveLossTask`](../contrastive_losses/ContrastiveLossTask.md) + +
+contrastive_losses.DeepGraphInfomaxTask(
+    *args, **kwargs
+)
+
+ + + + + + + + + + + + + + + + + + + + + + + + + + + +

Args

+node_set_name + +Name of the node set for readout. +
+feature_name + +Feature name for readout. +
+representations_layer_name + +Layer name for uncorrupted representations. +
+corruptor + +Corruptor instance for creating negative samples. If not +specified, we use ShuffleFeaturesGlobally by default. +
+projector_units + +Sequence of layer sizes for projector network. +Projectors prevent dimensional collapse, but can hinder training for +easy corruptions. For more details, see +https://arxiv.org/abs/2304.12210. +
+seed + +Random seed for the default corruptor (ShuffleFeaturesGlobally). +
+ +## Methods + +

losses

+ +View +source + +
+losses() -> runner.Losses
+
+ +Returns arbitrary task specific losses. + +

make_contrastive_layer

+ +View +source + +
+make_contrastive_layer() -> tf.keras.layers.Layer
+
+ +Returns the layer contrasting clean outputs with the correupted ones. + +

metrics

+ +View +source + +
+metrics() -> runner.Metrics
+
+ +Returns arbitrary task specific metrics. + +

predict

+ +View +source + +
+predict(
+    *args
+) -> runner.Predictions
+
+ +Apply a readout head for use with various contrastive losses. + + + + + + + + + + + +
Args
+*args + +A tuple of (clean, corrupted) tfgnn.GraphTensors. +
+ + + + + + + + + + +
Returns
+The logits for some contrastive loss as produced by the implementing +subclass. +
+ +

preprocess

+ +View +source + +
+preprocess(
+    inputs: GraphTensor
+) -> tuple[Sequence[GraphTensor], Mapping[str, Field]]
+
+ +Creates labels--i.e., (positive, negative)--for Deep Graph Infomax. diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DropoutFeatures.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DropoutFeatures.md new file mode 100644 index 00000000..161ecffb --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/DropoutFeatures.md @@ -0,0 +1,56 @@ +# contrastive_losses.DropoutFeatures + + + + + View source +on GitHub + +Base class for graph corruptor. + +Inherits From: [`Corruptor`](../contrastive_losses/Corruptor.md) + +
+contrastive_losses.DropoutFeatures(
+    *args, seed: Optional[float] = None, **kwargs
+)
+
+ + + + + + + + + + + + + + + + + + + + + +

Args

+corruption_spec + +A spec for corruption application. +
+corruption_fn + +Corruption function. +
+default + +Global application default of the corruptor. This is only used +when corruption_spec is None. +
+**kwargs + +Additional keyword arguments. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ShuffleFeaturesGlobally.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ShuffleFeaturesGlobally.md new file mode 100644 index 00000000..88c73fa3 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/ShuffleFeaturesGlobally.md @@ -0,0 +1,60 @@ +# contrastive_losses.ShuffleFeaturesGlobally + + + + + View source +on GitHub + +A corruptor that shuffles features. + +Inherits From: [`Corruptor`](../contrastive_losses/Corruptor.md) + +
+contrastive_losses.ShuffleFeaturesGlobally(
+    *args, seed: Optional[float] = None, **kwargs
+)
+
+ + + +NOTE: this function does not currently support TPUs. Consider using other +corruptor functions if executing on TPUs. See b/269249455 for reference. + + + + + + + + + + + + + + + + + + + + +

Args

+corruption_spec + +A spec for corruption application. +
+corruption_fn + +Corruption function. +
+default + +Global application default of the corruptor. This is only used +when corruption_spec is None. +
+**kwargs + +Additional keyword arguments. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletEmbeddingSquaredDistances.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletEmbeddingSquaredDistances.md new file mode 100644 index 00000000..7af68145 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletEmbeddingSquaredDistances.md @@ -0,0 +1,17 @@ +# contrastive_losses.TripletEmbeddingSquaredDistances + + + + + View source +on GitHub + +Computes embeddings distance between positive and negative pairs. + +
+contrastive_losses.TripletEmbeddingSquaredDistances(
+    trainable=True, name=None, dtype=None, dynamic=False, **kwargs
+)
+
+ + diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletLossTask.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletLossTask.md new file mode 100644 index 00000000..82e5af85 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/TripletLossTask.md @@ -0,0 +1,200 @@ +# contrastive_losses.TripletLossTask + + + + + View source +on GitHub + +The triplet loss task. + +Inherits From: +[`ContrastiveLossTask`](../contrastive_losses/ContrastiveLossTask.md) + +
+contrastive_losses.TripletLossTask(
+    *args, margin: float = 1.0, **kwargs
+)
+
+ + + + + + + + + + + + + + + + + + + + + + + + + + + +

Args

+node_set_name + +Name of the node set for readout. +
+feature_name + +Feature name for readout. +
+representations_layer_name + +Layer name for uncorrupted representations. +
+corruptor + +Corruptor instance for creating negative samples. If not +specified, we use ShuffleFeaturesGlobally by default. +
+projector_units + +Sequence of layer sizes for projector network. +Projectors prevent dimensional collapse, but can hinder training for +easy corruptions. For more details, see +https://arxiv.org/abs/2304.12210. +
+seed + +Random seed for the default corruptor (ShuffleFeaturesGlobally). +
+ +## Methods + +

losses

+ +View +source + +
+losses() -> runner.Losses
+
+ +Returns arbitrary task specific losses. + +

make_contrastive_layer

+ +View +source + +
+make_contrastive_layer() -> tf.keras.layers.Layer
+
+ +Returns the layer contrasting clean outputs with the correupted ones. + +

metrics

+ +View +source + +
+metrics() -> runner.Metrics
+
+ +Returns arbitrary task specific metrics. + +

predict

+ +View +source + +
+predict(
+    *args
+) -> runner.Predictions
+
+ +Apply a readout head for use with triplet contrastive loss. + + + + + + + + + + + +
Args
+*args + +A tuple of (anchor, positive_sample, negative_sample) +tfgnn.GraphTensors. +
+ + + + + + + + + + +
Returns
+The positive and negative distance embeddings for triplet loss as produced +by the implementing subclass. +
+ +

preprocess

+ +View +source + +
+preprocess(
+    inputs: GraphTensor
+) -> tuple[Sequence[GraphTensor], tfgnn.Field]
+
+ +Creates unused pseudo-labels. + +The input tensor should have the anchor and positive sample stacked along the +first dimension for each feature within each node set. The corruptor is applied +on the positive sample. + + + + + + + + + + + +
Args
+inputs + +The anchor and positive sample stack along the first axis. +
+ + + + + + + + + + +
Returns
+Sequence of three graph tensors (anchor, positive_sample, +corrupted_sample) and unused pseudo-labels. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/VicRegTask.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/VicRegTask.md new file mode 100644 index 00000000..cf466b46 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/VicRegTask.md @@ -0,0 +1,169 @@ +# contrastive_losses.VicRegTask + + + + + View source +on GitHub + +A VICReg Task. + +Inherits From: +[`ContrastiveLossTask`](../contrastive_losses/ContrastiveLossTask.md) + +
+contrastive_losses.VicRegTask(
+    *args,
+    sim_weight: Union[tf.Tensor, float] = 25.0,
+    var_weight: Union[tf.Tensor, float] = 25.0,
+    cov_weight: Union[tf.Tensor, float] = 1.0,
+    **kwargs
+)
+
+ + + + + + + + + + + + + + + + + + + + + + + + + + + +

Args

+node_set_name + +Name of the node set for readout. +
+feature_name + +Feature name for readout. +
+representations_layer_name + +Layer name for uncorrupted representations. +
+corruptor + +Corruptor instance for creating negative samples. If not +specified, we use ShuffleFeaturesGlobally by default. +
+projector_units + +Sequence of layer sizes for projector network. +Projectors prevent dimensional collapse, but can hinder training for +easy corruptions. For more details, see +https://arxiv.org/abs/2304.12210. +
+seed + +Random seed for the default corruptor (ShuffleFeaturesGlobally). +
+ +## Methods + +

losses

+ +View +source + +
+losses() -> runner.Losses
+
+ +Returns arbitrary task specific losses. + +

make_contrastive_layer

+ +View +source + +
+make_contrastive_layer() -> tf.keras.layers.Layer
+
+ +Returns the layer contrasting clean outputs with the correupted ones. + +

metrics

+ +View +source + +
+metrics() -> runner.Metrics
+
+ +Returns arbitrary task specific metrics. + +

predict

+ +View +source + +
+predict(
+    *args
+) -> runner.Predictions
+
+ +Apply a readout head for use with various contrastive losses. + + + + + + + + + + + +
Args
+*args + +A tuple of (clean, corrupted) tfgnn.GraphTensors. +
+ + + + + + + + + + +
Returns
+The logits for some contrastive loss as produced by the implementing +subclass. +
+ +

preprocess

+ +View +source + +
+preprocess(
+    inputs: GraphTensor
+) -> tuple[Sequence[GraphTensor], runner.Predictions]
+
+ +Applies a `Corruptor` and returns empty pseudo-labels. diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/all_symbols.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/all_symbols.md new file mode 100644 index 00000000..72b37858 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/all_symbols.md @@ -0,0 +1,24 @@ +# All symbols in TensorFlow GNN Models: fContrastiveLosses + + + +## Primary symbols + +* contrastive_losses +* contrastive_losses.AllSvdMetrics +* contrastive_losses.BarlowTwinsTask +* contrastive_losses.ContrastiveLossTask +* contrastive_losses.CorruptionSpec +* contrastive_losses.Corruptor +* contrastive_losses.DeepGraphInfomaxLogits +* contrastive_losses.DeepGraphInfomaxTask +* contrastive_losses.DropoutFeatures +* contrastive_losses.ShuffleFeaturesGlobally +* contrastive_losses.TripletEmbeddingSquaredDistances +* contrastive_losses.TripletLossTask +* contrastive_losses.VicRegTask +* contrastive_losses.coherence +* contrastive_losses.numerical_rank +* contrastive_losses.pseudo_condition_number +* contrastive_losses.rankme +* contrastive_losses.self_clustering diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/coherence.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/coherence.md new file mode 100644 index 00000000..66ffc367 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/coherence.md @@ -0,0 +1,69 @@ +# contrastive_losses.coherence + + + + + View source +on GitHub + +Coherence metric implementation. + +
+@tf.function
+contrastive_losses.coherence(
+    representations: tf.Tensor,
+    *,
+    sigma: Optional[tf.Tensor] = None,
+    u: Optional[tf.Tensor] = None
+) -> tf.Tensor
+
+ + + +Coherence measures how easy it is to construct a linear classifier on top of +data without knowing downstream labels. Refer to +https://arxiv.org/abs/2305.16562 for more details. + + + + + + + + + + + + + + + + + +

Args

+representations + +Input representations, a rank-2 tensor. +
+sigma + +Unused. +
+u + +An optional tensor with left singular vectors of representations. If not +present, computes a SVD of representations. +
+ + + + + + + + + + +

Returns

+Metric value as scalar tf.Tensor. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/numerical_rank.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/numerical_rank.md new file mode 100644 index 00000000..567971ef --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/numerical_rank.md @@ -0,0 +1,70 @@ +# contrastive_losses.numerical_rank + + + + + View source +on GitHub + +Numerical rank implementation. + +
+@tf.function
+contrastive_losses.numerical_rank(
+    representations: tf.Tensor,
+    *,
+    sigma: Optional[tf.Tensor] = None,
+    u: Optional[tf.Tensor] = None
+) -> tf.Tensor
+
+ + + +Computes a metric that estimates the numerical column rank of a matrix. Rank is +estimated as a matrix trace divided by the largest eigenvalue. When our matrix +is a covariance matrix, we can compute both the trace and the largest eigenvalue +efficiently via SVD as the largest singular value squared. + + + + + + + + + + + + + + + + + +

Args

+representations + +Input representations. We expect rank 2 input. +
+sigma + +An optional tensor with singular values of representations. If not +present, computes SVD (singular values only) of representations. +
+u + +Unused. +
+ + + + + + + + + + +

Returns

+Metric value as scalar tf.Tensor. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/pseudo_condition_number.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/pseudo_condition_number.md new file mode 100644 index 00000000..1ece1671 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/pseudo_condition_number.md @@ -0,0 +1,69 @@ +# contrastive_losses.pseudo_condition_number + + + + + View source +on GitHub + +Pseudo-condition number metric implementation. + +
+@tf.function
+contrastive_losses.pseudo_condition_number(
+    representations: tf.Tensor,
+    *,
+    sigma: Optional[tf.Tensor] = None,
+    u: Optional[tf.Tensor] = None
+) -> tf.Tensor
+
+ + + +Computes a metric that measures the decay rate of the singular values. NOTE: Can +be unstable in practice, when using small batch sizes, leading to numerical +instabilities. + + + + + + + + + + + + + + + + + +

Args

+representations + +Input representations. We expect rank 2 input. +
+sigma + +An optional tensor with singular values of representations. If not +present, computes SVD (singular values only) of representations. +
+u + +Unused. +
+ + + + + + + + + + +

Returns

+Metric value as scalar tf.Tensor. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/rankme.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/rankme.md new file mode 100644 index 00000000..90223259 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/rankme.md @@ -0,0 +1,77 @@ +# contrastive_losses.rankme + + + + + View source +on GitHub + +RankMe metric implementation. + +
+@tf.function
+contrastive_losses.rankme(
+    representations: tf.Tensor,
+    *,
+    sigma: Optional[tf.Tensor] = None,
+    u: Optional[tf.Tensor] = None,
+    epsilon: float = 1e-12,
+    **_
+) -> tf.Tensor
+
+ + + +Computes a metric that measures the decay rate of the singular values. For the +paper, see https://arxiv.org/abs/2210.02885. + + + + + + + + + + + + + + + + + + + + +

Args

+representations + +Input representations as rank-2 tensor. +
+sigma + +An optional tensor with singular values of representations. If not +present, computes SVD (singular values only) of representations. +
+u + +Unused. +
+epsilon + +Epsilon for numerican stability. +
+ + + + + + + + + + +

Returns

+Metric value as scalar tf.Tensor. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/self_clustering.md b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/self_clustering.md new file mode 100644 index 00000000..474189d5 --- /dev/null +++ b/tensorflow_gnn/docs/api_docs/python/models/contrastive_losses/self_clustering.md @@ -0,0 +1,63 @@ +# contrastive_losses.self_clustering + + + + + View source +on GitHub + +Self-clustering metric implementation. + +
+@tf.function
+contrastive_losses.self_clustering(
+    representations: tf.Tensor, *, subtract_mean: bool = False, **_
+) -> tf.Tensor
+
+ + + +Computes a metric that measures how well distributed representations are, if +projected on the unit sphere. If `subtract_mean` is True, we additionally remove +the mean from representations. The metric has a range of (-0.5, 1]. It achieves +its maximum of 1 if representations collapse to a single point, and it is +approximately 0 if representations are distributed randomly on the sphere. In +theory, it can achieve negative values if the points are maximally equiangular, +although this is very rare in practice. Refer to +https://arxiv.org/abs/2305.16562 for more details. + + + + + + + + + + + + + + +

Args

+representations + +Input representations. +
+subtract_mean + +Whether to subtract the mean from representations. +
+ + + + + + + + + + +

Returns

+Metric value as scalar tf.Tensor. +
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2.md index 6a8cf186..2292e9c9 100644 --- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2.md +++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2.md @@ -1,17 +1,10 @@ # Module: gat_v2 -[TOC] - - - - + + View source +on GitHub Graph Attention Networks v2. diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2Conv.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2Conv.md index 677ce5e8..361de99b 100644 --- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2Conv.md +++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2Conv.md @@ -1,17 +1,10 @@ # gat_v2.GATv2Conv -[TOC] - - - - + + View source +on GitHub The multi-head attention from Graph Attention Networks v2 (GATv2). @@ -86,77 +79,81 @@ example, if the input features have shape `[num_nodes, 2, 4, 1]`, then it will perform an identical computation on each of the `num_nodes * 2 * 4` input values. +This layer can be restored from config by `tf.keras.models.load_model()` when +saved as part of a Keras model using `save_format="tf"`. + +

Init args

-`num_heads` +num_heads The number of attention heads.
-`per_head_channels` +per_head_channels The number of channels for each attention head. This means: - if `heads_merge_type == "concat"`, then final output size will be: - `per_head_channels * num_heads`. - if `heads_merge_type == "mean"`, then final output size will be: - `per_head_channels`. + if heads_merge_type == "concat", then final output size will be: + per_head_channels * num_heads. + if heads_merge_type == "mean", then final output size will be: + per_head_channels.
-`receiver_tag` +receiver_tag -one of `tfgnn.SOURCE`, `tfgnn.TARGET` or `tfgnn.CONTEXT`. +one of tfgnn.SOURCE, tfgnn.TARGET or tfgnn.CONTEXT. The results of attention are aggregated for this graph piece. -If set to `tfgnn.SOURCE` or `tfgnn.TARGET`, the layer can be called for +If set to tfgnn.SOURCE or tfgnn.TARGET, the layer can be called for an edge set and will aggregate results at the specified endpoint of the edges. -If set to `tfgnn.CONTEXT`, the layer can be called for an edge set or +If set to tfgnn.CONTEXT, the layer can be called for an edge set or node set. If left unset for init, the tag must be passed at call time.
-`receiver_feature` +receiver_feature -Can be set to override `tfgnn.HIDDEN_STATE` for use as +Can be set to override tfgnn.HIDDEN_STATE for use as the receiver's input feature to attention. (The attention key is derived from this input.)
-`sender_node_feature` +sender_node_feature -Can be set to override `tfgnn.HIDDEN_STATE` for use as +Can be set to override tfgnn.HIDDEN_STATE for use as the input feature from sender nodes to attention. -IMPORTANT: Must be set to `None` for use with `receiver_tag=tfgnn.CONTEXT` +IMPORTANT: Must be set to None for use with receiver_tag=tfgnn.CONTEXT on an edge set, or for pooling from edges without sender node states.
-`sender_edge_feature` +sender_edge_feature Can be set to a feature name of the edge set to select -it as an input feature. By default, this set to `None`, which disables +it as an input feature. By default, this set to None, which disables this input. -IMPORTANT: Must be set for use with `receiver_tag=tfgnn.CONTEXT` +IMPORTANT: Must be set for use with receiver_tag=tfgnn.CONTEXT on an edge set.
-`use_bias` +use_bias If true, a bias term is added to the transformations of query and @@ -164,7 +161,7 @@ value inputs.
-`edge_dropout` +edge_dropout Can be set to a dropout rate for edge dropout. (When pooling @@ -173,29 +170,29 @@ is dropped out.)
-`attention_activation` +attention_activation The nonlinearity used on the transformed inputs before multiplying with the trained weights of the attention layer. This can be specified as a Keras layer, a tf.keras.activations.* -function, or a string understood by `tf.keras.layers.Activation()`. +function, or a string understood by tf.keras.layers.Activation(). Defaults to "leaky_relu", which in turn defaults to a negative slope -of `alpha=0.2`. +of alpha=0.2.
-`heads_merge_type` +heads_merge_type The merge operation for combining output from -all `num_heads` attention heads. By default, output of heads will be +all num_heads attention heads. By default, output of heads will be concatenated. However, GAT paper (Velickovic et al, Eq 6) recommends *only for output layer* to do mean across attention heads, which is acheivable -by setting to `"mean"`. +by setting to "mean".
-`activation` +activation The nonlinearity applied to the final result of attention, @@ -203,17 +200,17 @@ specified in the same ways as attention_activation.
-`kernel_initializer` +kernel_initializer -Can be set to a `kernel_initializer` as understood -by `tf.keras.layers.Dense` etc. -An `Initializer` object gets cloned before use to ensure a fresh seed, -if not set explicitly. For more, see `tfgnn.keras.clone_initializer()`. +Can be set to a kernel_initializer as understood +by tf.keras.layers.Dense etc. +An Initializer object gets cloned before use to ensure a fresh seed, +if not set explicitly. For more, see tfgnn.keras.clone_initializer().
-`kernel_regularizer` +kernel_regularizer If given, will be used to regularize all layer kernels. @@ -222,56 +219,58 @@ If given, will be used to regularize all layer kernels.
+ - + + +

Args

`receiver_tag` one of -`tfgnn.SOURCE`, `tfgnn.TARGET` or `tfgnn.CONTEXT`. The results are aggregated -for this graph piece. If set to `tfgnn.SOURCE` or `tfgnn.TARGET`, the layer can -be called for an edge set and will aggregate results at the specified endpoint -of the edges. If set to `tfgnn.CONTEXT`, the layer can be called for an edge set -or a node set and will aggregate results for context (per graph component). If -left unset for init, the tag must be passed at call time.
-`receiver_feature` The name of the -feature that is read from the receiver graph piece and passed as -convolve(receiver_input=...).
-`sender_node_feature` The name of the -feature that is read from the sender nodes, if any, and passed as -convolve(sender_node_input=...). NOTICE this must be `None` for use with -`receiver_tag=tfgnn.CONTEXT` on an edge set, or for pooling from edges without -sender node states.
-`sender_edge_feature` The name of the -feature that is read from the sender edges, if any, and passed as -convolve(sender_edge_input=...). NOTICE this must not be `None` for use with -`receiver_tag=tfgnn.CONTEXT` on an edge set.
-`extra_receiver_ops` A str-keyed -dictionary of Python callables that are wrapped to bind some arguments and then -passed on to `convolve()`. Sample usage: `extra_receiver_ops={"softmax": -tfgnn.softmax}`. The values passed in this dict must be callable as follows, -with two positional arguments: +
receiver_tag one of +tfgnn.SOURCE, tfgnn.TARGET or +tfgnn.CONTEXT. The results are aggregated for this graph piece. If +set to tfgnn.SOURCE or tfgnn.TARGET, the layer can be +called for an edge set and will aggregate results at the specified endpoint of +the edges. If set to tfgnn.CONTEXT, the layer can be called for an +edge set or a node set and will aggregate results for context (per graph +component). If left unset for init, the tag must be passed at call time.
receiver_feature The name of the feature that is read from the receiver graph piece and +passed as convolve(receiver_input=...).
+sender_node_feature The +name of the feature that is read from the sender nodes, if any, and passed as +convolve(sender_node_input=...). NOTICE this must be None for use +with receiver_tag=tfgnn.CONTEXT on an edge set, or for pooling from +edges without sender node states.
+sender_edge_feature The +name of the feature that is read from the sender edges, if any, and passed as +convolve(sender_edge_input=...). NOTICE this must not be None for +use with receiver_tag=tfgnn.CONTEXT on an edge set.
extra_receiver_ops A +str-keyed dictionary of Python callables that are wrapped to bind some arguments +and then passed on to convolve(). Sample usage: +extra_receiver_ops={"softmax": tfgnn.softmax}. The values passed in +this dict must be callable as follows, with two positional arguments: ```python f(graph, receiver_tag, node_set_name=..., feature_value=..., ...) f(graph, receiver_tag, edge_set_name=..., feature_value=..., ...) ``` -The wrapped callables seen by `convolve()` can be called like +The wrapped callables seen by convolve() can be called like ```python wrapped_f(feature_value, ...) ``` -The first three arguments of `f` are set to the input GraphTensor of -the layer and the tag/name pair required by `tfgnn.broadcast()` and -`tfgnn.pool()` to move values between the receiver and the messages that +The first three arguments of f are set to the input GraphTensor of +the layer and the tag/name pair required by tfgnn.broadcast() and +tfgnn.pool() to move values between the receiver and the messages that are computed inside the convolution. The sole positional argument of -`wrapped_f()` is passed to `f()` as `feature_value=`, and any keyword +wrapped_f() is passed to f() as feature_value=, and any keyword arguments are forwarded.
-`**kwargs` +**kwargs Forwarded to the base class tf.keras.layers.Layer. @@ -280,30 +279,31 @@ Forwarded to the base class tf.keras.layers.Layer.
+

Attributes

-`takes_receiver_input` +takes_receiver_input -If `False`, all calls to convolve() will get `receiver_input=None`. +If False, all calls to convolve() will get receiver_input=None.
-`takes_sender_edge_input` +takes_sender_edge_input -If `False`, all calls to convolve() will get `sender_edge_input=None`. +If False, all calls to convolve() will get sender_edge_input=None.
-`takes_sender_node_input` +takes_sender_node_input -If `False`, all calls to convolve() will get `sender_node_input=None`. +If False, all calls to convolve() will get sender_node_input=None.
@@ -312,7 +312,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.

convolve

-View +View source 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2EdgePool.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2EdgePool.md
index 9cd423a8..8b6c5c74 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2EdgePool.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2EdgePool.md
@@ -1,17 +1,10 @@
 # gat_v2.GATv2EdgePool
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a layer for pooling edges with GATv2-style attention.
 
@@ -41,21 +34,26 @@ an edge set to do the analogous pooling of edge states to context.
 NOTE: This layer cannot pool node states. For that, use
 gat_v2.GATv2Conv.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  
 
@@ -64,11 +67,11 @@ A new `GATv2MPNNGraphUpdate` object.
 
 

 
 
Args

 
 

 
-`num_heads`
+num_heads
 
 
 The number of attention heads.
 
 

 
-`per_head_channels`
+per_head_channels
 
 
 The number of channels for each attention head. This
@@ -63,20 +61,20 @@ means that the final output size will be per_head_channels * num_heads.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
 The results of attention are aggregated for this graph piece.
-If set to `tfgnn.CONTEXT`, the layer can be called for an edge set or
+If set to tfgnn.CONTEXT, the layer can be called for an edge set or
 node set.
-If set to an IncidentNodeTag (e.g., `tfgnn.SOURCE` or `tfgnn.TARGET`),
+If set to an IncidentNodeTag (e.g., tfgnn.SOURCE or tfgnn.TARGET),
 the layer can be called for an edge set and will aggregate results at
 the specified endpoint of the edges.
 If left unset, the tag must be passed when calling the layer.
 
 

 
-`receiver_feature`
+receiver_feature
 
 
 By default, the default state feature of the receiver
@@ -85,7 +83,7 @@ selected by setting this argument.
 
 

 
-`sender_feature`
+sender_feature
 
 
 By default, the default state feature of the edge set is
@@ -94,7 +92,7 @@ selected by setting this argument.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Any other option for GATv2Conv, except sender_node_feature,
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2HomGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2HomGraphUpdate.md
index cbc66046..acc03341 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2HomGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2HomGraphUpdate.md
@@ -1,17 +1,10 @@
 # gat_v2.GATv2HomGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a GraphUpdate layer with a Graph Attention Network V2 (GATv2).
 
@@ -40,21 +33,26 @@ objects instead, such as the GATv2MPNNGraphUpdate.
 > explicitly stored in the input GraphTensor. Attention of a node to itself
 > requires having an explicit loop in the edge set.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  
+
+
+

 
 
Args

 
 

 
-`num_heads`
+num_heads
 
 
 The number of attention heads.
 
 

 
-`per_head_channels`
+per_head_channels
 
 
 The number of channels for each attention head. This
@@ -62,22 +60,22 @@ means that the final output size will be per_head_channels * num_heads.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-one of `tfgnn.SOURCE` or `tfgnn.TARGET`.
+one of tfgnn.SOURCE or tfgnn.TARGET.
 
 

 
-`feature_name`
+feature_name
 
 
 The feature name of node states; defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`heads_merge_type`
+heads_merge_type
 
 
 "concat" or "mean". Gets passed to GATv2Conv, which uses
@@ -85,14 +83,14 @@ it to combine all heads into layer's output.
 
 

 
-`name`
+name
 
 
 Optionally, a name for the layer returned.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Any optional arguments to GATv2Conv, see there.
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2MPNNGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2MPNNGraphUpdate.md
index 0b67e958..ad5a214e 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2MPNNGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/GATv2MPNNGraphUpdate.md
@@ -1,17 +1,10 @@
 # gat_v2.GATv2MPNNGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a GraphUpdate layer for message passing with GATv2 pooling.
 
@@ -44,37 +37,42 @@ the messages and their pooling with attention, followed by a dense layer to
 compute the new node states from a concatenation of the old node state and all
 pooled messages.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`units`
+units
 
 
 The dimension of output hidden states for each node.
 
 

 
-`message_dim`
+message_dim
 
 
 The dimension of messages (attention values) computed on
-each edge.  Must be divisible by `num_heads`.
+each edge.  Must be divisible by num_heads.
 
 

 
-`num_heads`
+num_heads
 
 
-The number of attention heads used by GATv2. `message_dim`
+The number of attention heads used by GATv2. message_dim
 must be divisible by this number.
 
 

 
-`heads_merge_type`
+heads_merge_type
 
 
 "concat" or "mean". Gets passed to GATv2Conv, which uses
@@ -82,15 +80,15 @@ it to combine all heads into layer's output.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-one of `tfgnn.TARGET` or `tfgnn.SOURCE`, to select the
+one of tfgnn.TARGET or tfgnn.SOURCE, to select the
 incident node of each edge that receives the message.
 
 

 
-`node_set_names`
+node_set_names
 
 
 The names of node sets to update. If unset, updates all
@@ -98,16 +96,16 @@ that are on the receiving end of any edge set.
 
 

 
-`edge_feature`
+edge_feature
 
 
 Can be set to a feature name of the edge set to select
-it as an input feature. By default, this set to `None`, which disables
+it as an input feature. By default, this set to None, which disables
 this input.
 
 

 
-`l2_regularization`
+l2_regularization
 
 
 The coefficient of L2 regularization for weights and
@@ -115,7 +113,7 @@ biases.
 
 

 
-`edge_dropout_rate`
+edge_dropout_rate
 
 
 The edge dropout rate applied during attention pooling
@@ -123,26 +121,26 @@ of edges.
 
 

 
-`state_dropout_rate`
+state_dropout_rate
 
 
 The dropout rate applied to the resulting node states.
 
 

 
-`attention_activation`
+attention_activation
 
 
 The nonlinearity used on the transformed inputs
 before multiplying with the trained weights of the attention layer.
 This can be specified as a Keras layer, a tf.keras.activations.*
-function, or a string understood by `tf.keras.layers.Activation()`.
+function, or a string understood by tf.keras.layers.Activation().
 Defaults to "leaky_relu", which in turn defaults to a negative slope
-of `alpha=0.2`.
+of alpha=0.2.
 
 

 
-`conv_activation`
+conv_activation
 
 
 The nonlinearity applied to the result of attention on one
@@ -150,7 +148,7 @@ edge set, specified in the same ways as attention_activation.
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the new node states computed by
@@ -158,23 +156,24 @@ this graph update.
 
 

 
-`kernel_initializer`
+kernel_initializer
 
 
-Can be set to a `kernel_initializer` as understood
-by `tf.keras.layers.Dense` etc.
+Can be set to a kernel_initializer as understood
+by tf.keras.layers.Dense etc.
 
 

 

 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_from_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_from_config_dict.md
index 858724c5..ef2f3c62 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_from_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_from_config_dict.md
@@ -1,23 +1,18 @@
 # gat_v2.graph_update_from_config_dict
 
-[TOC]
-
 
 
-

 
 
Returns

 

 
 A GraphUpdate layer for use on a scalar GraphTensor with
-`tfgnn.HIDDEN_STATE` features on the node sets.
+tfgnn.HIDDEN_STATE features on the node sets.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns a GATv2MPNNGraphUpdate initialized from `cfg`.
 
 
 gat_v2.graph_update_from_config_dict(
-    cfg: config_dict.ConfigDict
+    cfg: config_dict.ConfigDict,
+    *,
+    node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None
 ) -> tf.keras.layers.Layer
 

 
@@ -30,15 +25,23 @@ Returns a GATv2MPNNGraphUpdate initialized from `cfg`.
 
 

 
-`cfg`
+cfg
 
 
-A `ConfigDict` with the fields defined by
-`graph_update_get_config_dict()`. All fields with non-`None` values are
+A ConfigDict with the fields defined by
+graph_update_get_config_dict(). All fields with non-None values are
 used as keyword arguments for initializing and returning a
-`GATv2MPNNGraphUpdate` object. For the required arguments of
-`GATv2MPNNGraphUpdate.__init__`, users must set a value in
-`cfg` before passing it here.
+GATv2MPNNGraphUpdate object. For the required arguments of
+GATv2MPNNGraphUpdate.__init__, users must set a value in
+cfg before passing it here.
+

+node_set_names
+
+Optionally, the names of NodeSets to update; forwarded to
+GATv2MPNNGraphUpdate.__init__.
 
 

 

@@ -50,7 +53,7 @@ used as keyword arguments for initializing and returning a
 
Returns

 

 
-A new `GATv2MPNNGraphUpdate` object.
+A new GATv2MPNNGraphUpdate object.
 
 

 
-`TypeError`
+TypeError
 
 
-if `cfg` fails to supply a required argument for
-`GATv2MPNNGraphUpdate.__init__`.
+if cfg fails to supply a required argument for
+GATv2MPNNGraphUpdate.__init__.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_get_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_get_config_dict.md
index 0f48b98f..073dfac2 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_get_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gat_v2/graph_update_get_config_dict.md
@@ -1,17 +1,10 @@
 # gat_v2.graph_update_get_config_dict
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns ConfigDict for graph_update_from_config_dict() with defaults.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gcn.md b/tensorflow_gnn/docs/api_docs/python/models/gcn.md
index 0ecf4de6..103f9866 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gcn.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gcn.md
@@ -1,17 +1,10 @@
 # Module: gcn
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Graph Convolutional Networks.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNConv.md b/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNConv.md
index d2890c55..ddbeeb2c 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNConv.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNConv.md
@@ -1,17 +1,10 @@
 # gcn.GCNConv
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Implements the Graph Convolutional Network by Kipf&Welling (2016).
 
@@ -71,14 +64,18 @@ $$v_{ij} = w_{ij} / (\sqrt{\deg^{in}_i} \sqrt{\deg^{in}_j}).$$
 For symmetric graphs (as in the original GCN paper), `"in_out"` and `"in_in"`
 are equal, but the latter needs to compute degrees just once.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Init arguments

 
 

 
-`units`
+units
 
 
 Number of output units for this transformation applied to sender
@@ -86,26 +83,26 @@ node features.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
 This layer's result is obtained by pooling the per-edge
-results at this endpoint of each edge. The default is `tfgnn.TARGET`,
+results at this endpoint of each edge. The default is tfgnn.TARGET,
 but it is perfectly reasonable to do a convolution towards the
-`tfgnn.SOURCE` instead. (Source and target are conventional names for
+tfgnn.SOURCE instead. (Source and target are conventional names for
 the incident nodes of a directed edge, data flow in a GNN may happen
 in either direction.)
 
 

 
-`activation`
+activation
 
 
 Keras activation to apply to the result, defaults to 'relu'.
 
 

 
-`use_bias`
+use_bias
 
 
 Whether to add bias in the final transformation. The original
@@ -114,7 +111,7 @@ with Keras and other implementations.
 
 

 
-`add_self_loops`
+add_self_loops
 
 
 Whether to compute the result as if a loop from each node
@@ -123,24 +120,24 @@ with an edge weight of one.
 
 

 
-`kernel_initializer`
+kernel_initializer
 
 
-Can be set to a `kernel_initializer` as understood
-by `tf.keras.layers.Dense` etc.
-An `Initializer` object gets cloned before use to ensure a fresh seed,
-if not set explicitly. For more, see `tfgnn.keras.clone_initializer()`.
+Can be set to a kernel_initializer as understood
+by tf.keras.layers.Dense etc.
+An Initializer object gets cloned before use to ensure a fresh seed,
+if not set explicitly. For more, see tfgnn.keras.clone_initializer().
 
 

 
-`node_feature`
+node_feature
 
 
 Name of the node feature to transform.
 
 

 
-`edge_weight_feature_name`
+edge_weight_feature_name
 
 
 Can be set to the name of a feature on the edge
@@ -149,15 +146,15 @@ it as the edge's entry in the adjacency matrix, instead of the default 1.
 
 

 
-`degree_normalization`
+degree_normalization
 
 
-Can be set to `"none"`, `"in"`, `"out"`, `"in_out"`,
-or `"in_in"`, as explained above.
+Can be set to "none", "in", "out", "in_out",
+or "in_in", as explained above.
 
 

 
-`**kwargs`
+**kwargs
 
 
 additional arguments for the Layer.
@@ -166,23 +163,24 @@ additional arguments for the Layer.
 

 
 
+
  

 
 
Call arguments

 
 

 
-`graph`
+graph
 
 
 The GraphTensor on which to apply the layer.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
-Edge set of `graph` over which to apply the layer.
+Edge set of graph over which to apply the layer.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNHomGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNHomGraphUpdate.md
index 1d3767e2..3b4663a6 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNHomGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNHomGraphUpdate.md
@@ -1,17 +1,10 @@
 # gcn.GCNHomGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a graph update layer for GCN convolution.
 
@@ -41,32 +34,37 @@ GCNConv objects instead.
 > requires having an explicit loop in the edge set, or setting
 > `add_self_loops=True`.
 
+Thie layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`units`
+units
 
 
 The dimension of output hidden states for each node.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-The default is `tfgnn.TARGET`,
+The default is tfgnn.TARGET,
 but it is perfectly reasonable to do a convolution towards the
-`tfgnn.SOURCE` instead. (Source and target are conventional names for
+tfgnn.SOURCE instead. (Source and target are conventional names for
 the incident nodes of a directed edge, data flow in a GNN may happen
 in either direction.)
 
 

 
-`add_self_loops`
+add_self_loops
 
 
 Whether to compute the result as if a loop from each node
@@ -74,22 +72,22 @@ to itself had been added to the edge set.
 
 

 
-`feature_name`
+feature_name
 
 
 The feature name of node states; defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`name`
+name
 
 
 Optionally, a name for the layer returned.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Any optional arguments to GCNConv, see there.
diff --git a/tensorflow_gnn/docs/api_docs/python/models/graph_sage.md b/tensorflow_gnn/docs/api_docs/python/models/graph_sage.md
index e418b9a1..09192446 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/graph_sage.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/graph_sage.md
@@ -1,17 +1,10 @@
 # Module: graph_sage
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 GraphSAGE.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GCNGraphSAGENodeSetUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GCNGraphSAGENodeSetUpdate.md
index dc2240b4..415ce1a5 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GCNGraphSAGENodeSetUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GCNGraphSAGENodeSetUpdate.md
@@ -1,17 +1,10 @@
 # graph_sage.GCNGraphSAGENodeSetUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 GCNGraphSAGENodeSetUpdate is an extension of the mean aggregator operator.
 
@@ -77,33 +70,37 @@ reduce operation, instead only the sender node states will be accumulated based
 on the reduce_type specified. If share_weights is set to True, then single
 weight matrix will be used in place of W_E and W_self.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`edge_set_names`
+edge_set_names
 
 
 A list of edge set names to broadcast sender node states.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-Either one of `tfgnn.SOURCE` or `tfgnn.TARGET`. The results
+Either one of tfgnn.SOURCE or tfgnn.TARGET. The results
 of GraphSAGE convolution are aggregated for this graph piece. If set to
-`tfgnn.SOURCE` or `tfgnn.TARGET`, the layer will be called for each edge
+tfgnn.SOURCE or tfgnn.TARGET, the layer will be called for each edge
 set and will aggregate results at the specified endpoint of the edges.
 This should point at the node_set_name for each of the specified edge
 set name in the edge_set_name_dict.
 
 

 
-`reduce_type`
+reduce_type
 
 
 An aggregation operation name. Supported list of aggregation
@@ -111,24 +108,24 @@ operators are sum or mean.
 
 

 
-`self_node_feature`
+self_node_feature
 
 
 Feature name for the self node sets to be aggregated
 with the broadcasted sender node states. Default is
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`sender_node_feature`
+sender_node_feature
 
 
 Feature name for the sender node sets. Default is
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`units`
+units
 
 
 Number of output units for the linear transformation applied to
@@ -136,7 +133,7 @@ sender node and self node features.
 
 

 
-`dropout_rate`
+dropout_rate
 
 
 Can be set to a dropout rate that will be applied to both
@@ -144,7 +141,7 @@ self node and the sender node states.
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the update node states. This can
@@ -153,7 +150,7 @@ string understood by tf.keras.layers.Activation(). Defaults to relu.
 
 

 
-`use_bias`
+use_bias
 
 
 If true a bias term will be added to mean aggregated feature
@@ -161,7 +158,7 @@ vectors before applying non-linear activation.
 
 

 
-`share_weights`
+share_weights
 
 
 If left unset, separate weights are used to transform the
@@ -171,7 +168,7 @@ applied to all inputs.
 
 

 
-`add_self_loop`
+add_self_loop
 
 
 If left at True (the default), each node state update takes
diff --git a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEAggregatorConv.md b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEAggregatorConv.md
index 3bf2a5cd..83cc06ac 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEAggregatorConv.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEAggregatorConv.md
@@ -1,17 +1,10 @@
 # graph_sage.GraphSAGEAggregatorConv
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 GraphSAGE: element-wise aggregation of neighbors and their linear
 transformation.
@@ -45,42 +38,46 @@ besides "mean", see the reduce_type=... argument. For stateful transformation
 with a hidden layer, see
 graph_sage.GraphSAGEPoolingConv.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`receiver_tag`
+receiver_tag
 
 
-Either one of `tfgnn.SOURCE` or `tfgnn.TARGET`. The results
+Either one of tfgnn.SOURCE or tfgnn.TARGET. The results
 of GraphSAGE convolution are aggregated for this graph piece. If set to
-`tfgnn.SOURCE` or `tfgnn.TARGET`, the layer will be called for an edge
+tfgnn.SOURCE or tfgnn.TARGET, the layer will be called for an edge
 set and will aggregate results at the specified endpoint of the edges.
 
 

 
-`reduce_type`
+reduce_type
 
 
 An aggregation operation name. Supported list of aggregation
 operators can be found at
-`tfgnn.get_registered_reduce_operation_names()`.
+tfgnn.get_registered_reduce_operation_names().
 
 

 
-`sender_node_feature`
+sender_node_feature
 
 
 Can be set to specify the feature name for use as the
 input feature from sender nodes to GraphSAGE aggregation, defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`units`
+units
 
 
 Number of output units for the linear transformation applied to
@@ -88,7 +85,7 @@ sender node features.
 
 

 
-`dropout_rate`
+dropout_rate
 
 
 Can be set to a dropout rate that will be applied to sender
@@ -96,7 +93,7 @@ node features (independently on each edge).
 
 

 
-`**kwargs`
+**kwargs
 
 
 Additional arguments for the Layer.
@@ -105,30 +102,31 @@ Additional arguments for the Layer.
 

 
 
+
  

 
 
Attributes

 
 

 
-`takes_receiver_input`
+takes_receiver_input
 
 
-If `False`, all calls to convolve() will get `receiver_input=None`.
+If False, all calls to convolve() will get receiver_input=None.
 
 

 
-`takes_sender_edge_input`
+takes_sender_edge_input
 
 
-If `False`, all calls to convolve() will get `sender_edge_input=None`.
+If False, all calls to convolve() will get sender_edge_input=None.
 
 

 
-`takes_sender_node_input`
+takes_sender_node_input
 
 
-If `False`, all calls to convolve() will get `sender_node_input=None`.
+If False, all calls to convolve() will get sender_node_input=None.
 
 

 

@@ -137,7 +135,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 
convolve

 
-View
+View
 source
 
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEGraphUpdate.md
index 3a3fb4c3..c7156448 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEGraphUpdate.md
@@ -1,17 +1,10 @@
 # graph_sage.GraphSAGEGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a GraphSAGE GraphUpdater layer for nodes in node_set_names.
 
@@ -43,13 +36,14 @@ applies only one step of GraphSAGE convolution over the incident nodes of the
 edge_set_name_list for the specified node_set_name node.
 
 
+
  
 
@@ -64,11 +66,11 @@ A new Layer object as returned by `MtAlbisGraphUpdate()`.
 
 

 
 
Args

 
 

 
-`units`
+units
 
 
 Number of output units of the linear transformation applied to both
@@ -57,7 +51,7 @@ final aggregated sender node features as well as the self node feature.
 
 

 
-`hidden_units`
+hidden_units
 
 
 Number of output units to be configure for GraphSAGE pooling
@@ -65,17 +59,17 @@ type convolution only.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-Either one of `tfgnn.SOURCE` or `tfgnn.TARGET`. The results of
-GraphSAGE are aggregated for this graph piece. When set to `tfgnn.SOURCE`
-or `tfgnn.TARGET`, the layer is called for an edge set and will aggregate
+Either one of tfgnn.SOURCE or tfgnn.TARGET. The results of
+GraphSAGE are aggregated for this graph piece. When set to tfgnn.SOURCE
+or tfgnn.TARGET, the layer is called for an edge set and will aggregate
 results at the specified endpoint of the edges.
 
 

 
-`node_set_names`
+node_set_names
 
 
 By default, this layer updates all node sets that receive
@@ -86,15 +80,15 @@ auxiliary node sets.
 
 

 
-`reduce_type`
+reduce_type
 
 
 An aggregation operation name. Supported list of aggregation
-operators can be found at `tfgnn.get_registered_reduce_operation_names()`.
+operators can be found at tfgnn.get_registered_reduce_operation_names().
 
 

 
-`use_pooling`
+use_pooling
 
 
 If enabled, graph_sage.GraphSAGEPoolingConv will be used,
@@ -103,7 +97,7 @@ provided edges.
 
 

 
-`use_bias`
+use_bias
 
 
 If true a bias term will be added to the linear transformations
@@ -111,7 +105,7 @@ for the incident node features as well as for the self node feature.
 
 

 
-`dropout_rate`
+dropout_rate
 
 
 Can be set to a dropout rate that will be applied to both
@@ -119,7 +113,7 @@ incident node features as well as the self node feature.
 
 

 
-`l2_normalize`
+l2_normalize
 
 
 If enabled l2 normalization will be applied to final node
@@ -127,7 +121,7 @@ states.
 
 

 
-`combine_type`
+combine_type
 
 
 Can be set to "sum" or "concat". If it's specified as concat
@@ -136,32 +130,32 @@ node state will be added with the sender node features.
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the concatenated or added node state
 and aggregated sender node features. This can be specified as a Keras
 layer, a tf.keras.activations.* function, or a string understood by
-`tf.keras.layers.Activation()`. Defaults to relu.
+tf.keras.layers.Activation(). Defaults to relu.
 
 

 
-`feature_name`
+feature_name
 
 
 The feature name of node states; defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`name`
+name
 
 
 Optionally, a name for the layer returned.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Any optional arguments to graph_sage.GraphSAGEPoolingConv,
diff --git a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGENextState.md b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGENextState.md
index 62e06194..68d0c6c8 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGENextState.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGENextState.md
@@ -1,17 +1,10 @@
 # graph_sage.GraphSAGENextState
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 GraphSAGENextState: compute new node states with GraphSAGE algorithm.
 
@@ -60,7 +53,7 @@ equal, unless `combine_type="concat"`is set.
 GraphSAGE is Algorithm 1 in Hamilton et al.:
 ["Inductive Representation Learning on Large Graphs"](https://arxiv.org/abs/1706.02216),
 2017. It computes the new hidden state h_v for each node v from a concatenation
-of the previous hiddden state with an aggregation of the neighbor states as
+of the previous hidden state with an aggregation of the neighbor states as
 
 $$h_v = \sigma(W \text{ concat}(h_v, h_{N(v)}))$$
 
@@ -87,13 +80,14 @@ Beyond the original GraphSAGE, this class supports:
 *   additional options to influence normalization, activation, etc.
 
 
+
  
+
+
+

 
 
Args

 
 

 
-`units`
+units
 
 
 Number of output units for the linear transformation applied to the
@@ -101,7 +95,7 @@ node feature.
 
 

 
-`use_bias`
+use_bias
 
 
 If true a bias term will be added to the linear transformations
@@ -109,7 +103,7 @@ for the self node feature.
 
 

 
-`dropout_rate`
+dropout_rate
 
 
 Can be set to a dropout rate that will be applied to the
@@ -117,15 +111,15 @@ node feature.
 
 

 
-`feature_name`
+feature_name
 
 
 The feature name of node states; defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`l2_normalize`
+l2_normalize
 
 
 If enabled l2 normalization will be applied to node state
@@ -133,7 +127,7 @@ vectors.
 
 

 
-`combine_type`
+combine_type
 
 
 Can be set to "sum" or "concat". The default "sum" recovers
@@ -144,17 +138,17 @@ Setting this to "concat" concatenates the results of the transformations
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the concatenated or added node
 state and aggregated sender node features. This can be specified as a
 Keras layer, a tf.keras.activations.* function, or a string understood
-by `tf.keras.layers.Activation()`. Defaults to relu.
+by tf.keras.layers.Activation(). Defaults to relu.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Forwarded to the base class tf.keras.layers.Layer.
diff --git a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEPoolingConv.md b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEPoolingConv.md
index 3d53c792..a23368ed 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEPoolingConv.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/graph_sage/GraphSAGEPoolingConv.md
@@ -1,17 +1,10 @@
 # graph_sage.GraphSAGEPoolingConv
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 GraphSAGE: pooling aggregator transform of neighbors followed by linear
 transformation.
@@ -53,33 +46,37 @@ involves the aforementioned hidden layer. For element-wise aggregation (as in
 `tfgnn.pool_edges_to_node()`), see
 graph_sage.GraphSAGEAggregatorConv.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`receiver_tag`
+receiver_tag
 
 
-Either one of `tfgnn.SOURCE` or `tfgnn.TARGET`. The results
+Either one of tfgnn.SOURCE or tfgnn.TARGET. The results
 of GraphSAGE are aggregated for this graph piece. If set to
-`tfgnn.SOURCE` or `tfgnn.TARGET`, the layer will be called for an edge
+tfgnn.SOURCE or tfgnn.TARGET, the layer will be called for an edge
 set and will aggregate results at the specified endpoint of the edges.
 
 

 
-`sender_node_feature`
+sender_node_feature
 
 
 Can be set to specify the feature name for use as the
 input feature from sender nodes to GraphSAGE aggregation, defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`units`
+units
 
 
 Number of output units for the final dimensionality of the output
@@ -87,7 +84,7 @@ from the layer.
 
 

 
-`hidden_units`
+hidden_units
 
 
 Number of output units for the linear transformation applied
@@ -97,16 +94,16 @@ W_pool from Eq. (3) in
 
 

 
-`reduce_type`
+reduce_type
 
 
 An aggregation operation name. Supported list of aggregation
 operators can be found at
-`tfgnn.get_registered_reduce_operation_names()`.
+tfgnn.get_registered_reduce_operation_names().
 
 

 
-`use_bias`
+use_bias
 
 
 If true a bias term will be added to the linear transformations
@@ -114,7 +111,7 @@ for the sender node features.
 
 

 
-`dropout_rate`
+dropout_rate
 
 
 Can be set to a dropout rate that will be applied to sender
@@ -122,17 +119,17 @@ node features (independently on each edge).
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the concatenated or added node
 state and aggregated sender node features. This can be specified as a
 Keras layer, a tf.keras.activations.* function, or a string understood
-by `tf.keras.layers.Activation()`. Defaults to relu.
+by tf.keras.layers.Activation(). Defaults to relu.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Additional arguments for the Layer.
@@ -141,30 +138,31 @@ Additional arguments for the Layer.
 

 
 
+
  

 
 
Attributes

 
 

 
-`takes_receiver_input`
+takes_receiver_input
 
 
-If `False`, all calls to convolve() will get `receiver_input=None`.
+If False, all calls to convolve() will get receiver_input=None.
 
 

 
-`takes_sender_edge_input`
+takes_sender_edge_input
 
 
-If `False`, all calls to convolve() will get `sender_edge_input=None`.
+If False, all calls to convolve() will get sender_edge_input=None.
 
 

 
-`takes_sender_node_input`
+takes_sender_node_input
 
 
-If `False`, all calls to convolve() will get `sender_node_input=None`.
+If False, all calls to convolve() will get sender_node_input=None.
 
 

 

@@ -173,7 +171,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 
convolve

 
-View
+View
 source
 
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/mt_albis.md b/tensorflow_gnn/docs/api_docs/python/models/mt_albis.md
index 61451923..299ee813 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/mt_albis.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/mt_albis.md
@@ -1,17 +1,10 @@
 # Module: mt_albis
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 TF-GNN's Model Template "Albis".
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/mt_albis/MtAlbisGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/mt_albis/MtAlbisGraphUpdate.md
index 1b9c26f4..bfb58361 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/mt_albis/MtAlbisGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/mt_albis/MtAlbisGraphUpdate.md
@@ -1,17 +1,10 @@
 # mt_albis.MtAlbisGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns GraphUpdate layer for message passing with Model Template "Albis".
 
@@ -44,6 +37,10 @@ Returns GraphUpdate layer for message passing with Model Template "Albis".
 The TF-GNN Model Template "Albis" provides a small selection of field-tested GNN
 architectures through the unified interface of this class.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
 
  
@@ -52,29 +49,29 @@ architectures through the unified interface of this class.
 
 

 
-`units`
+units
 
 
 The dimension of node states in the output GraphTensor.
 
 

 
-`message_dim`
+message_dim
 
 
 The dimension of messages computed transiently on each edge.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-One of `tfgnn.SOURCE` or `tfgnn.TARGET`. The messages are
+One of tfgnn.SOURCE or tfgnn.TARGET. The messages are
 sent to the nodes at this endpoint of edges.
 
 

 
-`node_set_names`
+node_set_names
 
 
 Optionally, the names of NodeSets to update. By default,
@@ -82,7 +79,7 @@ all NodeSets are updated that receive from at least one EdgeSet.
 
 

 
-`edge_feature_name`
+edge_feature_name
 
 
 Optionally, the name of an edge feature to include in
@@ -90,53 +87,54 @@ message computation on edges.
 
 

 
-`attention_type`
+attention_type
 
 
-`"none"`, `"multi_head"`, or `"gat_v2"`. Selects whether
+"none", "multi_head", or "gat_v2". Selects whether
 messages are pooled with data-dependent weights computed by a trained
 attention mechansim.
 
 

 
-`attention_edge_set_names`
+attention_edge_set_names
 
 
 If set, edge sets other than those named here
-will be treated as if `attention_type="none"` regardless.
+will be treated as if attention_type="none" regardless.
 
 

 
-`attention_num_heads`
+attention_num_heads
 
 
-For attention_types `"multi_head"` or `"gat_v2"`,
+For attention_types "multi_head" or "gat_v2",
 the number of attention heads.
 
 

 
-`simple_conv_reduce_type`
+simple_conv_reduce_type
 
 
-For attention_type `"none"`, controls how messages
-are aggregated on an EdgeSet for each receiver node. Defaults to `"mean"`;
-other recommened values are the concatenations `"mean|sum"`, `"mean|max"`,
-and `"mean|sum|max"` (but mind the increased output dimension and the
-corresponding increase in the number of weights in the next-state layer).
-Technically, can be set to any reduce_type understood by `tfgnn.pool()`.
+For attention_type "none", controls how messages
+are aggregated on an EdgeSet for each receiver node. Defaults to "mean";
+other recommended values are the concatenations "mean|sum",
+"mean|max", and "mean|sum|max" (but mind the increased output
+dimension and the corresponding increase in the number of weights in the
+next-state layer). Technically, can be set to any reduce_type understood
+by tfgnn.pool().
 
 

 
-`simple_conv_use_receiver_state`
+simple_conv_use_receiver_state
 
 
-For attention_type `"none"`, controls
+For attention_type "none", controls
 whether the receiver node state is used in computing each edge's message
-(in addition to the sender node state and possibly an `edge feature`).
+(in addition to the sender node state and possibly an edge feature).
 
 

 
-`state_dropout_rate`
+state_dropout_rate
 
 
 The dropout rate applied to the pooled and combined
@@ -147,7 +145,7 @@ is applied to messages after pooling.)
 
 

 
-`edge_dropout_rate`
+edge_dropout_rate
 
 
 Can be set to a dropout rate for entire edges during
@@ -156,7 +154,7 @@ an edge is dropped, as if the edge were not present in the graph.
 
 

 
-`l2_regularization`
+l2_regularization
 
 
 The coefficient of L2 regularization for trained weights.
@@ -164,54 +162,54 @@ The coefficient of L2 regularization for trained weights.
 
 

 
-`kernel_initializer`
+kernel_initializer
 
 
-Can be set to a `kernel_initializer` as understood
-by `tf.keras.layers.Dense` etc.
-An `Initializer` object gets cloned before use to ensure a fresh seed,
-if not set explicitly. For more, see `tfgnn.keras.clone_initializer()`.
+Can be set to a kernel_initializer as understood
+by tf.keras.layers.Dense etc.
+An Initializer object gets cloned before use to ensure a fresh seed,
+if not set explicitly. For more, see tfgnn.keras.clone_initializer().
 
 

 
-`normalization_type`
+normalization_type
 
 
 controls the normalization of output node states.
-By default (`"layer"`), LayerNormalization is used. Can be set to
-`"none"`, or to `"batch"` for BatchNormalization.
+By default ("layer"), LayerNormalization is used. Can be set to
+"none", or to "batch" for BatchNormalization.
 
 

 
-`batch_normalization_momentum`
+batch_normalization_momentum
 
 
-If `normalization_type="batch"`, sets the
-`BatchNormalization(momentum=...)` parameter. Ignored otherwise.
+If normalization_type="batch", sets the
+BatchNormalization(momentum=...) parameter. Ignored otherwise.
 
 

 
-`next_state_type`
+next_state_type
 
 
-`"dense"` or `"residual"`. With the latter, a residual
+"dense" or "residual". With the latter, a residual
 link is added from the old to the new node state, which requires that all
-input node states already have size `units` (unless their size is 0, as
+input node states already have size units (unless their size is 0, as
 for latent node sets, in which case the residual link is omitted).
 
 

 
-`edge_set_combine_type`
+edge_set_combine_type
 
 
-`"concat"` or `"sum"`. Controls how pooled messages
+"concat" or "sum". Controls how pooled messages
 from various edge sets are combined as inputs to the NextState layer
-that updates the node states. Defaults to `"concat"`, which gives the
+that updates the node states. Defaults to "concat", which gives the
 pooled messages from each edge set separate weights in the NextState
-layer, namely `units * message_dim * num_incident_edge_sets` per node set.
-Setting this to `"sum"` adds up the pooled messages into a single
+layer, namely units * message_dim * num_incident_edge_sets per node set.
+Setting this to "sum" adds up the pooled messages into a single
 vector before passing them into the NextState layer, which requires just
-`units * message_dim` weights per node set.
+units * message_dim weights per node set.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_from_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_from_config_dict.md
index 6843881a..5eb5a0e6 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_from_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_from_config_dict.md
@@ -1,23 +1,17 @@
 # mt_albis.graph_update_from_config_dict
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Constructs a MtAlbisGraphUpdate from a ConfigDict.
 
 
 mt_albis.graph_update_from_config_dict(
-    cfg: config_dict.ConfigDict
+    cfg: config_dict.ConfigDict,
+    node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None
 ) -> tf.keras.layers.Layer
 

 
@@ -30,16 +24,24 @@ Constructs a MtAlbisGraphUpdate from a ConfigDict.
 
 

 
-`cfg`
+cfg
 
 
-A `ConfigDict` with the fields defined by
-`graph_update_get_config_dict()`. All fields with non-`None` values are
+A ConfigDict with the fields defined by
+graph_update_get_config_dict(). All fields with non-None values are
 used as keyword arguments for initializing and returning a
-`MtAlbisGraphUpdate` object. For the required arguments of
-`MtAlbisGraphUpdate.__init__`, users must set a value in `cfg` before
+MtAlbisGraphUpdate object. For the required arguments of
+MtAlbisGraphUpdate.__init__, users must set a value in cfg before
 passing it here.
 

+node_set_names
+
+Optionally, the names of NodeSets to update; forwarded to
+MtAlbisGraphUpdate.__init__.
+
 

 

 
@@ -50,7 +52,7 @@ passing it here.
 
Returns

 

 
-A new Layer object as returned by `MtAlbisGraphUpdate()`.
+A new Layer object as returned by MtAlbisGraphUpdate().
 
 

 
-`TypeError`
+TypeError
 
 
-if `cfg` fails to supply a required argument for
-`MtAlbisGraphUpdate()`.
+if cfg fails to supply a required argument for
+MtAlbisGraphUpdate().
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_get_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_get_config_dict.md
index a104fb31..4211a131 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_get_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/mt_albis/graph_update_get_config_dict.md
@@ -1,17 +1,10 @@
 # mt_albis.graph_update_get_config_dict
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns ConfigDict for graph_update_from_config_dict() with defaults.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention.md
index 47ae18eb..f4dc2f07 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention.md
@@ -1,17 +1,10 @@
 # Module: multi_head_attention
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Transformer-style multi-head attention.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionConv.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionConv.md
index 36aab5a5..c5d90833 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionConv.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionConv.md
@@ -1,17 +1,10 @@
 # multi_head_attention.MultiHeadAttentionConv
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Transformer-style (dot-product) multi-head attention on GNNs.
 
@@ -32,7 +25,7 @@ Transformer-style (dot-product) multi-head attention on GNNs.
     kernel_initializer: Any = None,
     kernel_regularizer: Any = None,
     transform_keys: bool = True,
-    score_scaling: Literal['none', 'rsqrt_dim', 'trainable_sigmoid'] = 'rsqrt_dim',
+    score_scaling: Literal['none', 'rsqrt_dim', 'trainable_elup1'] = 'rsqrt_dim',
     transform_values_after_pooling: bool = False,
     **kwargs
 )
@@ -96,50 +89,54 @@ using the other arguments.
 
 Example: Transformer-style attention on neighbors along incoming edges whose
 result is concatenated with the old node state and passed through a Dense layer
-to compute the new node state. `dense = tf.keras.layers.Dense graph =
-tfgnn.keras.layers.GraphUpdate( node_sets={"paper":
-tfgnn.keras.layers.NodeSetUpdate( {"cites":
-tfgnn.keras.layers.MultiHeadAttentionConv( message_dim,
-receiver_tag=tfgnn.TARGET)},
-tfgnn.keras.layers.NextStateFromConcat(dense(node_state_dim)))} )(graph)`
+to compute the new node state.
+
+```
+dense = tf.keras.layers.Dense
+graph = tfgnn.keras.layers.GraphUpdate(
+    node_sets={"paper": tfgnn.keras.layers.NodeSetUpdate(
+        {"cites": tfgnn.keras.layers.MultiHeadAttentionConv(
+             message_dim, receiver_tag=tfgnn.TARGET)},
+        tfgnn.keras.layers.NextStateFromConcat(dense(node_state_dim)))}
+)(graph)
+```
 
 For now, there is a variant that modifies the inputs transformation part and
 could potentially be beneficial:
 
-```
-1. (transform_keys is False) Instead of projecting both queries and
-  keys when computing attention weights, we only project the queries
-  because the two linear projections can be collapsed to a single
-  projection:
-
-    $$ (Q_v W_Q^k)(K_u W_K^k)^T
-      = Q_v (W_Q^k {W_K^k}^T) K_u^T
-      = Q_v W_{QK}^k K_u^T $$
-
-  where $d$ is the key width. (Following "Attention is all you need",
-  this scaling is meant to achieve unit variance of the results, assuming
-  that $Q_v W_{QK}^k$ has unit variance due to the initialization of
-  $Q_v W_{QK}^k$.)
-
-  NOTE: The single projection matrix behaves differently in
-  gradient-descent training than the product of two matrices.
-```
+1.  (transform_keys is False) Instead of projecting both queries and keys when
+    computing attention weights, we only project the queries because the two
+    linear projections can be collapsed to a single projection:
+
+    $$ (Q_v W_Q^k)(K_u W_K^k)^T = Q_v (W_Q^k {W_K^k}^T) K_u^T = Q_v W_{QK}^k
+    K_u^T $$
+
+    where $d$ is the key width. (Following "Attention is all you need", this
+    scaling is meant to achieve unit variance of the results, assuming that $Q_v
+    W_{QK}^k$ has unit variance due to the initialization of $Q_v W_{QK}^k$.)
+
+    NOTE: The single projection matrix behaves differently in gradient-descent
+    training than the product of two matrices.
+
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
 
 
+
  

 
 
Init args

 
 

 
-`num_heads`
+num_heads
 
 
 The number of attention heads.
 
 

 
-`per_head_channels`
+per_head_channels
 
 
 The number of channels for each attention head. This
@@ -147,51 +144,51 @@ means that the final output size will be per_head_channels * num_heads.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-one of `tfgnn.SOURCE`, `tfgnn.TARGET` or `tfgnn.CONTEXT`.
+one of tfgnn.SOURCE, tfgnn.TARGET or tfgnn.CONTEXT.
 The results of attention are aggregated for this graph piece.
-If set to `tfgnn.SOURCE` or `tfgnn.TARGET`, the layer can be called for
+If set to tfgnn.SOURCE or tfgnn.TARGET, the layer can be called for
 an edge set and will aggregate results at the specified endpoint of the
 edges.
-If set to `tfgnn.CONTEXT`, the layer can be called for an edge set or
+If set to tfgnn.CONTEXT, the layer can be called for an edge set or
 node set.
 If left unset for init, the tag must be passed at call time.
 
 

 
-`receiver_feature`
+receiver_feature
 
 
-Can be set to override `tfgnn.HIDDEN_STATE`
+Can be set to override tfgnn.HIDDEN_STATE
 for use as the receiver's input feature to attention. (The attention key
 is derived from this input.)
 
 

 
-`sender_node_feature`
+sender_node_feature
 
 
-Can be set to override `tfgnn.HIDDEN_STATE`
+Can be set to override tfgnn.HIDDEN_STATE
 for use as the input feature from sender nodes to attention.
-IMPORTANT: Must be set to `None` for use with `receiver_tag=tfgnn.CONTEXT`
+IMPORTANT: Must be set to None for use with receiver_tag=tfgnn.CONTEXT
 on an edge set, or for pooling from edges without sender node states.
 
 

 
-`sender_edge_feature`
+sender_edge_feature
 
 
 Can be set to a feature name of the edge set to select
-it as an input feature. By default, this set to `None`, which disables
+it as an input feature. By default, this set to None, which disables
 this input.
-IMPORTANT: Must be set for use with `receiver_tag=tfgnn.CONTEXT`
+IMPORTANT: Must be set for use with receiver_tag=tfgnn.CONTEXT
 on an edge set.
 
 

 
-`use_bias`
+use_bias
 
 
 If true, bias terms are added to the transformations of query,
@@ -199,7 +196,7 @@ key and value inputs.
 
 

 
-`edge_dropout`
+edge_dropout
 
 
 Can be set to a dropout rate for edge dropout. (When pooling
@@ -208,7 +205,7 @@ is dropped out.)
 
 

 
-`inputs_dropout`
+inputs_dropout
 
 
 Dropout rate for random dropout on the inputs to this
@@ -216,18 +213,18 @@ convolution layer, i.e. the receiver, sender node, and sender edge inputs.
 
 

 
-`attention_activation`
+attention_activation
 
 
 The nonlinearity used on the transformed inputs
-(query, and keys if `transform_keys` is `True`) before computing the
+(query, and keys if transform_keys is True) before computing the
 attention scores. This can be specified as a Keras layer, a
 tf.keras.activations.* function, or a string understood by
-`tf.keras.layers.Activation`. Defaults to None.
+tf.keras.layers.Activation. Defaults to None.
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the final result of attention,
@@ -235,25 +232,25 @@ specified in the same ways as attention_activation.
 
 

 
-`kernel_initializer`
+kernel_initializer
 
 
-Can be set to a `kernel_initializer` as understood
-by `tf.keras.layers.Dense` etc.
-An `Initializer` object gets cloned before use to ensure a fresh seed,
-if not set explicitly. For more, see `tfgnn.keras.clone_initializer()`.
+Can be set to a kernel_initializer as understood
+by tf.keras.layers.Dense etc.
+An Initializer object gets cloned before use to ensure a fresh seed,
+if not set explicitly. For more, see tfgnn.keras.clone_initializer().
 
 

 
-`kernel_regularizer`
+kernel_regularizer
 
 
-Can be set to a `kernel_regularized` as understood
-by `tf.keras.layers.Dense` etc.
+Can be set to a kernel_regularized as understood
+by tf.keras.layers.Dense etc.
 
 

 
-`transform_keys`
+transform_keys
 
 
 If true, transform both queries and keys inputs. Otherwise,
@@ -263,22 +260,24 @@ independent of this arg.)
 
 

 
-`score_scaling`
+score_scaling
 
 
-One of either `"none"`, `"rsqrt_dim"`, or
-`"trainable_sigmoid"`. If set to `"rsqrt_dim"`, the attention scores are
+One of either "rsqrt_dim" (default), "trainable_elup1",
+or "none". If set to "rsqrt_dim", the attention scores are
 divided by the square root of the dimension of keys (i.e.,
-`per_head_channels` if `transform_keys=True`, otherwise whatever the
-dimension of combined sender inputs is). If set to `"trainable_sigmoid"`,
-the scores are scaled with `sigmoid(x)`, where `x` is a trainable weight
-of the model that is initialized to `-5.0`, which initially makes all the
-attention weights equal and slowly ramps up as the other weights in the
-layer converge. Defaults to `"rsqrt_dim"`.
+per_head_channels if transform_keys=True, otherwise whatever the
+dimension of combined sender inputs is). If set to "trainable_elup1",
+the scores are scaled with elu(x) + 1, where elu is the Exponential
+Linear Unit (see tf.keras.activations.elu), and x is a per-head
+trainable weight of the model that is initialized to 0.0. Recall that
+elu(x) + 1 == exp(x) if x<0 else x+1, so the
+initial scaling factor is 1.0, decreases exponentially below 1.0, and
+grows linearly above 1.0.
 
 

 
-`transform_values_after_pooling`
+transform_values_after_pooling
 
 
 By default, each attention head applies
@@ -292,56 +291,58 @@ IMPORTANT: Toggling this option breaks checkpoint compatibility.
 

 
 
+
  
 
-
+
+
+

 
 
Args
  `receiver_tag`   one of
-`tfgnn.SOURCE`, `tfgnn.TARGET` or `tfgnn.CONTEXT`. The results are aggregated
-for this graph piece. If set to `tfgnn.SOURCE` or `tfgnn.TARGET`, the layer can
-be called for an edge set and will aggregate results at the specified endpoint
-of the edges. If set to `tfgnn.CONTEXT`, the layer can be called for an edge set
-or a node set and will aggregate results for context (per graph component). If
-left unset for init, the tag must be passed at call time.  
 
-`receiver_feature`   The name of the
-feature that is read from the receiver graph piece and passed as
-convolve(receiver_input=...).  
 
-`sender_node_feature`   The name of the
-feature that is read from the sender nodes, if any, and passed as
-convolve(sender_node_input=...). NOTICE this must be `None` for use with
-`receiver_tag=tfgnn.CONTEXT` on an edge set, or for pooling from edges without
-sender node states.  
 
-`sender_edge_feature`   The name of the
-feature that is read from the sender edges, if any, and passed as
-convolve(sender_edge_input=...). NOTICE this must not be `None` for use with
-`receiver_tag=tfgnn.CONTEXT` on an edge set.  
 
-`extra_receiver_ops`   A str-keyed
-dictionary of Python callables that are wrapped to bind some arguments and then
-passed on to `convolve()`. Sample usage: `extra_receiver_ops={"softmax":
-tfgnn.softmax}`. The values passed in this dict must be callable as follows,
-with two positional arguments:
+
  receiver_tag   one of
+tfgnn.SOURCE, tfgnn.TARGET or
+tfgnn.CONTEXT. The results are aggregated for this graph piece. If
+set to tfgnn.SOURCE or tfgnn.TARGET, the layer can be
+called for an edge set and will aggregate results at the specified endpoint of
+the edges. If set to tfgnn.CONTEXT, the layer can be called for an
+edge set or a node set and will aggregate results for context (per graph
+component). If left unset for init, the tag must be passed at call time. 
  receiver_feature  The name of the feature that is read from the receiver graph piece and
+passed as convolve(receiver_input=...).  
 
+sender_node_feature   The
+name of the feature that is read from the sender nodes, if any, and passed as
+convolve(sender_node_input=...). NOTICE this must be None for use
+with receiver_tag=tfgnn.CONTEXT on an edge set, or for pooling from
+edges without sender node states.  
 
+sender_edge_feature   The
+name of the feature that is read from the sender edges, if any, and passed as
+convolve(sender_edge_input=...). NOTICE this must not be None for
+use with receiver_tag=tfgnn.CONTEXT on an edge set.  
 extra_receiver_ops   A
+str-keyed dictionary of Python callables that are wrapped to bind some arguments
+and then passed on to convolve(). Sample usage:
+extra_receiver_ops={"softmax": tfgnn.softmax}. The values passed in
+this dict must be callable as follows, with two positional arguments:
 
 ```python
 f(graph, receiver_tag, node_set_name=..., feature_value=..., ...)
 f(graph, receiver_tag, edge_set_name=..., feature_value=..., ...)
 ```
 
-The wrapped callables seen by `convolve()` can be called like
+The wrapped callables seen by convolve() can be called like
 
 ```python
 wrapped_f(feature_value, ...)
 ```
 
-The first three arguments of `f` are set to the input GraphTensor of
-the layer and the tag/name pair required by `tfgnn.broadcast()` and
-`tfgnn.pool()` to move values between the receiver and the messages that
+The first three arguments of f are set to the input GraphTensor of
+the layer and the tag/name pair required by tfgnn.broadcast() and
+tfgnn.pool() to move values between the receiver and the messages that
 are computed inside the convolution. The sole positional argument of
-`wrapped_f()` is passed to `f()`  as `feature_value=`, and any keyword
+wrapped_f() is passed to f()  as feature_value=, and any keyword
 arguments are forwarded.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Forwarded to the base class tf.keras.layers.Layer.
@@ -350,30 +351,31 @@ Forwarded to the base class tf.keras.layers.Layer.
 

 
 
+
  

 
 
Attributes

 
 

 
-`takes_receiver_input`
+takes_receiver_input
 
 
-If `False`, all calls to convolve() will get `receiver_input=None`.
+If False, all calls to convolve() will get receiver_input=None.
 
 

 
-`takes_sender_edge_input`
+takes_sender_edge_input
 
 
-If `False`, all calls to convolve() will get `sender_edge_input=None`.
+If False, all calls to convolve() will get sender_edge_input=None.
 
 

 
-`takes_sender_node_input`
+takes_sender_node_input
 
 
-If `False`, all calls to convolve() will get `sender_node_input=None`.
+If False, all calls to convolve() will get sender_node_input=None.
 
 

 

@@ -382,7 +384,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 
convolve

 
-View
+View
 source
 
 
@@ -408,44 +410,45 @@ from nodes to context). In the end, values have to be pooled from there into a
 Tensor with a leading dimension indexed by receivers, see `pool_to_receiver`.
 
 
+
  

 
 
Args

 
 

 
-`sender_node_input`
+sender_node_input
 
 
-The input Tensor from the sender NodeSet, or `None`.
-If self.takes_sender_node_input is `False`, this arg will be `None`.
-(If it is `True`, that depends on how this layer gets called.)
+The input Tensor from the sender NodeSet, or None.
+If self.takes_sender_node_input is False, this arg will be None.
+(If it is True, that depends on how this layer gets called.)
 See also broadcast_from_sender_node.
 
 

 
-`sender_edge_input`
+sender_edge_input
 
 
-The input Tensor from the sender EdgeSet, or `None`.
-If self.takes_sender_edge_input is `False`, this arg will be `None`.
-(If it is `True`, it depends on how this layer gets called.)
+The input Tensor from the sender EdgeSet, or None.
+If self.takes_sender_edge_input is False, this arg will be None.
+(If it is True, it depends on how this layer gets called.)
 If present, this Tensor is already indexed by the items for which
 messages are computed.
 
 

 
-`receiver_input`
+receiver_input
 
 
 The input Tensor from the receiver NodeSet or Context,
-or None. If self.takes_receiver_input is `False`, this arg will be
-`None`. (If it is `True`, it depends on how this layer gets called.)
+or None. If self.takes_receiver_input is False, this arg will be
+None. (If it is True, it depends on how this layer gets called.)
 See broadcast_from_receiver.
 
 

 
-`broadcast_from_sender_node`
+broadcast_from_sender_node
 
 
 A function that broadcasts a Tensor indexed
@@ -454,25 +457,25 @@ messages are computed.
 
 

 
-`broadcast_from_receiver`
+broadcast_from_receiver
 
 
-Call this as `broadcast_from_receiver(value)`
+Call this as broadcast_from_receiver(value)
 to broadcast a Tensor indexed like receiver_input to a Tensor indexed
 by the items for which messages are computed.
 
 

 
-`pool_to_receiver`
+pool_to_receiver
 
 
-Call this as `pool_to_receiver(value, reduce_type=...)`
+Call this as pool_to_receiver(value, reduce_type=...)
 to pool an item-indexed Tensor to a receiver-indexed tensor, using
 a reduce_type understood by tfgnn.pool(), such as "sum".
 
 

 
-`extra_receiver_ops`
+extra_receiver_ops
 
 
 The extra_receiver_ops passed to init, see there,
@@ -482,10 +485,10 @@ this argument, so subclass implementors not using it can omit it.
 
 

 
-`training`
+training
 
 
-The `training` boolean that was passed to Layer.call(). If true,
+The training boolean that was passed to Layer.call(). If true,
 the result is computed for training rather than inference. For example,
 calls to tf.nn.dropout() are usually conditioned on this flag.
 By contrast, calling another Keras layer (like tf.keras.layers.Dropout)
@@ -495,6 +498,7 @@ does not require forwarding this arg, Keras does that automatically.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionEdgePool.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionEdgePool.md
index d535a182..6fa90816 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionEdgePool.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionEdgePool.md
@@ -1,17 +1,10 @@
 # multi_head_attention.MultiHeadAttentionEdgePool
 
-[TOC]
-
 
 
-

 
 
Returns

-
-
+
+ View source
+on GitHub 
 
 Returns a layer for pooling edges with Transformer-style Multi-Head Attention.
 
@@ -40,21 +33,26 @@ an edge set to do the analogous pooling of edge states to context.
 
 NOTE: This layer cannot pool node states. For that, use MultiHeadAttentionConv.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`num_heads`
+num_heads
 
 
 The number of attention heads.
 
 

 
-`per_head_channels`
+per_head_channels
 
 
 The number of channels for each attention head. This
@@ -62,19 +60,19 @@ means that the final output size will be per_head_channels * num_heads.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
 The results of attention are aggregated for this graph piece.
-If set to `tfgnn.CONTEXT`, the layer can be called for an edge set or node
-set. If set to an IncidentNodeTag (e.g., `tfgnn.SOURCE` or
-`tfgnn.TARGET`), the layer can be called for an edge set and will
+If set to tfgnn.CONTEXT, the layer can be called for an edge set or node
+set. If set to an IncidentNodeTag (e.g., tfgnn.SOURCE or
+tfgnn.TARGET), the layer can be called for an edge set and will
 aggregate results at the specified endpoint of the edges. If left unset,
 the tag must be passed when calling the layer.
 
 

 
-`receiver_feature`
+receiver_feature
 
 
 By default, the default state feature of the receiver is
@@ -83,7 +81,7 @@ selected by setting this argument.
 
 

 
-`sender_feature`
+sender_feature
 
 
 By default, the default state feature of the edge set is
@@ -92,7 +90,7 @@ selected by setting this argument.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Any other option for MultiHeadAttentionConv, except
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionHomGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionHomGraphUpdate.md
index 27cc34b8..d532de39 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionHomGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionHomGraphUpdate.md
@@ -1,17 +1,10 @@
 # multi_head_attention.MultiHeadAttentionHomGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a GraphUpdate layer with a transformer-style multihead attention.
 
@@ -42,21 +35,26 @@ details).
 > that are explicitly stored in the input GraphTensor. Attention of a node to
 > itself requires having an explicit loop in the edge set.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`num_heads`
+num_heads
 
 
 The number of attention heads.
 
 

 
-`per_head_channels`
+per_head_channels
 
 
 The number of channels for each attention head. This
@@ -64,29 +62,29 @@ means that the final output size will be per_head_channels * num_heads.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-one of `tfgnn.SOURCE` or `tfgnn.TARGET`.
+one of tfgnn.SOURCE or tfgnn.TARGET.
 
 

 
-`feature_name`
+feature_name
 
 
 The feature name of node states; defaults to
-`tfgnn.HIDDEN_STATE`.
+tfgnn.HIDDEN_STATE.
 
 

 
-`name`
+name
 
 
 Optionally, a name for the layer returned.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Any optional arguments to MultiHeadAttentionConv, see there.
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionMPNNGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionMPNNGraphUpdate.md
index d6eb153e..9c496104 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionMPNNGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/MultiHeadAttentionMPNNGraphUpdate.md
@@ -1,17 +1,10 @@
 # multi_head_attention.MultiHeadAttentionMPNNGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a GraphUpdate layer for message passing with MultiHeadAttention pooling.
 
@@ -44,45 +37,50 @@ layer to compute the new node states from a concatenation of the old node state
 and all pooled messages, analogous to TF-GNN's
 `vanilla_mpnn.VanillaMPNNGraphUpdate` and `gat_v2.GATv2MPNNGraphUpdate`.
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`units`
+units
 
 
 The dimension of output hidden states for each node.
 
 

 
-`message_dim`
+message_dim
 
 
 The dimension of messages (attention values) computed on each
-edge.  Must be divisible by `num_heads`.
+edge.  Must be divisible by num_heads.
 
 

 
-`num_heads`
+num_heads
 
 
 The number of attention heads used by MultiHeadAttention.
-`message_dim` must be divisible by this number.
+message_dim must be divisible by this number.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-one of `tfgnn.TARGET` or `tfgnn.SOURCE`, to select the
+one of tfgnn.TARGET or tfgnn.SOURCE, to select the
 incident node of each edge that receives the message.
 
 

 
-`node_set_names`
+node_set_names
 
 
 The names of node sets to update. If unset, updates all that
@@ -90,16 +88,16 @@ are on the receiving end of any edge set.
 
 

 
-`edge_feature`
+edge_feature
 
 
 Can be set to a feature name of the edge set to select it as
-an input feature. By default, this set to `None`, which disables this
+an input feature. By default, this set to None, which disables this
 input.
 
 

 
-`l2_regularization`
+l2_regularization
 
 
 The coefficient of L2 regularization for weights and
@@ -107,7 +105,7 @@ biases.
 
 

 
-`edge_dropout_rate`
+edge_dropout_rate
 
 
 The edge dropout rate applied during attention pooling of
@@ -115,24 +113,24 @@ edges.
 
 

 
-`state_dropout_rate`
+state_dropout_rate
 
 
 The dropout rate applied to the resulting node states.
 
 

 
-`attention_activation`
+attention_activation
 
 
 The nonlinearity used on the transformed inputs before
 multiplying with the trained weights of the attention layer. This can be
 specified as a Keras layer, a tf.keras.activations.* function, or a string
-understood by `tf.keras.layers.Activation`. Defaults to None.
+understood by tf.keras.layers.Activation. Defaults to None.
 
 

 
-`conv_activation`
+conv_activation
 
 
 The nonlinearity applied to the result of attention on one
@@ -140,7 +138,7 @@ edge set, specified in the same ways as attention_activation.
 
 

 
-`activation`
+activation
 
 
 The nonlinearity applied to the new node states computed by this
@@ -148,25 +146,26 @@ graph update.
 
 

 
-`kernel_initializer`
+kernel_initializer
 
 
-Can be set to a `kernel_initializer` as understood
-by `tf.keras.layers.Dense` etc.
-An `Initializer` object gets cloned before use to ensure a fresh seed,
-if not set explicitly. For more, see `tfgnn.keras.clone_initializer()`.
+Can be set to a kernel_initializer as understood
+by tf.keras.layers.Dense etc.
+An Initializer object gets cloned before use to ensure a fresh seed,
+if not set explicitly. For more, see tfgnn.keras.clone_initializer().
 
 

 

 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_from_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_from_config_dict.md
index f6bbf78e..efac65b5 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_from_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_from_config_dict.md
@@ -1,72 +1,77 @@
 # multi_head_attention.graph_update_from_config_dict
 
-[TOC]
-
 
 
-

 
 
Returns

 

 
 A GraphUpdate layer for use on a scalar GraphTensor with
-`tfgnn.HIDDEN_STATE` features on the node sets.
+tfgnn.HIDDEN_STATE features on the node sets.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns a MultiHeadAttentionMPNNGraphUpdate initialized from `cfg`.
 
 
 multi_head_attention.graph_update_from_config_dict(
-    cfg: config_dict.ConfigDict
+    cfg: config_dict.ConfigDict,
+    *,
+    node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None
 ) -> tf.keras.layers.Layer
 

 
 
-
 
+
  
+
+
+

 
 
Args

 
 

 
-`cfg`
+cfg
 
 
-A `ConfigDict` with the fields defined by
-`graph_update_get_config_dict()`. All fields with non-`None` values are
+A ConfigDict with the fields defined by
+graph_update_get_config_dict(). All fields with non-None values are
 used as keyword arguments for initializing and returning a
-`MultiHeadAttentionMPNNGraphUpdate` object. For the required arguments of
-`MultiHeadAttentionMPNNGraphUpdate.__init__`, users must set a value in
-`cfg` before passing it here.
+MultiHeadAttentionMPNNGraphUpdate object. For the required arguments of
+MultiHeadAttentionMPNNGraphUpdate.__init__, users must set a value in
+cfg before passing it here.
+

+node_set_names
+
+Optionally, the names of NodeSets to update; forwarded to
+MultiHeadAttentionMPNNGraphUpdate.__init__.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A new `MultiHeadAttentionMPNNGraphUpdate` object.
+A new MultiHeadAttentionMPNNGraphUpdate object.
 
 

 
 

 
 
+
  

 
 
Raises

 
 

 
-`TypeError`
+TypeError
 
 
-if `cfg` fails to supply a required argument for
-`MultiHeadAttentionMPNNGraphUpdate.__init__`.
+if cfg fails to supply a required argument for
+MultiHeadAttentionMPNNGraphUpdate.__init__.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_get_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_get_config_dict.md
index 017fba0a..51a8aee9 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_get_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_get_config_dict.md
@@ -1,17 +1,10 @@
 # multi_head_attention.graph_update_get_config_dict
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns ConfigDict for graph_update_from_config_dict() with defaults.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn.md b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn.md
index 755c6346..148b9f45 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn.md
@@ -1,17 +1,10 @@
 # Module: vanilla_mpnn
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 TF-GNN's "Vanilla MPNN" model.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/VanillaMPNNGraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/VanillaMPNNGraphUpdate.md
index 00b4f949..9159630a 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/VanillaMPNNGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/VanillaMPNNGraphUpdate.md
@@ -1,17 +1,10 @@
 # vanilla_mpnn.VanillaMPNNGraphUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns a GraphUpdate layer for a Vanilla MPNN.
 
@@ -55,36 +48,41 @@ and the pooled messages from all incident node sets E_1, E_2, ...:
 $$h_v := \text{ReLU}(
     W_{\text{state}} (h_v || m_{E_1} || m_{E_2} || \ldots)).$$
 
+The layer returned by this function can be restored from config by
+`tf.keras.models.load_model()` when saved as part of a Keras model using
+`save_format="tf"`.
+
 
+
  

 
 
Args

 
 

 
-`units`
+units
 
 
 The dimension of output hidden states for each node.
 
 

 
-`message_dim`
+message_dim
 
 
 The dimension of messages computed on each edge.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
-one of `tfgnn.TARGET` or `tfgnn.SOURCE`, to select the
+one of tfgnn.TARGET or tfgnn.SOURCE, to select the
 incident node of each edge that receives the message.
 
 

 
-`node_set_names`
+node_set_names
 
 
 The names of node sets to update. If unset, updates all
@@ -92,26 +90,26 @@ that are on the receiving end of any edge set.
 
 

 
-`edge_feature`
+edge_feature
 
 
 Can be set to a feature name of the edge set to select
-it as an input feature. By default, this set to `None`, which disables
+it as an input feature. By default, this set to None, which disables
 this input.
 
 

 
-`reduce_type`
+reduce_type
 
 
 How to pool the messages from edges to receiver nodes; defaults
-to `"sum"`. Can be any reduce_type understood by `tfgnn.pool()`, including
-concatenations like `"sum|max"` (but mind the increased dimension of the
+to "sum". Can be any reduce_type understood by tfgnn.pool(), including
+concatenations like "sum|max" (but mind the increased dimension of the
 result and the growing number of model weights in the next-state layer).
 
 

 
-`l2_regularization`
+l2_regularization
 
 
 The coefficient of L2 regularization for weights and
@@ -119,7 +117,7 @@ biases.
 
 

 
-`dropout_rate`
+dropout_rate
 
 
 The dropout rate applied to messages on each edge and to the
@@ -127,17 +125,17 @@ new node state.
 
 

 
-`kernel_initializer`
+kernel_initializer
 
 
-Can be set to a `kernel_initializer` as understood
-by `tf.keras.layers.Dense` etc.
-An `Initializer` object gets cloned before use to ensure a fresh seed,
-if not set explicitly. For more, see `tfgnn.keras.clone_initializer()`.
+Can be set to a kernel_initializer as understood
+by tf.keras.layers.Dense etc.
+An Initializer object gets cloned before use to ensure a fresh seed,
+if not set explicitly. For more, see tfgnn.keras.clone_initializer().
 
 

 
-`use_layer_normalization`
+use_layer_normalization
 
 
 Flag to determine whether to apply layer
@@ -147,13 +145,14 @@ normalization to the new node state.
 

 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_from_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_from_config_dict.md
index 539deb0f..c08e7a78 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_from_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_from_config_dict.md
@@ -1,72 +1,77 @@
 # vanilla_mpnn.graph_update_from_config_dict
 
-[TOC]
-
 
 
-

 
 
Returns

 

 
 A GraphUpdate layer for use on a scalar GraphTensor with
-`tfgnn.HIDDEN_STATE` features on the node sets.
+tfgnn.HIDDEN_STATE features on the node sets.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns a VanillaMPNNGraphUpdate initialized from `cfg`.
 
 
 vanilla_mpnn.graph_update_from_config_dict(
-    cfg: config_dict.ConfigDict
+    cfg: config_dict.ConfigDict,
+    *,
+    node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None
 ) -> tf.keras.layers.Layer
 

 
 
-
 
+
  
+
+
+

 
 
Args

 
 

 
-`cfg`
+cfg
 
 
-A `ConfigDict` with the fields defined by
-`graph_update_get_config_dict()`. All fields with non-`None` values are
+A ConfigDict with the fields defined by
+graph_update_get_config_dict(). All fields with non-None values are
 used as keyword arguments for initializing and returning a
-`VanillaMPNNGraphUpdate` object. For the required arguments of
-`VanillaMPNNGraphUpdate.__init__`, users must set a value in `cfg` before
+VanillaMPNNGraphUpdate object. For the required arguments of
+VanillaMPNNGraphUpdate.__init__, users must set a value in cfg before
 passing it here.
 

+node_set_names
+
+Optionally, the names of NodeSets to update; forwarded to
+MtAlbisGraphUpdate.__init__.
+
 

 

 
 
+
  

 
 
Returns

 

 
-A new `VanillaMPNNGraphUpdate` object.
+A new VanillaMPNNGraphUpdate object.
 
 

 
 

 
 
+
  

 
 
Raises

 
 

 
-`TypeError`
+TypeError
 
 
-if `cfg` fails to supply a required argument for
-`VanillaMPNNGraphUpdate.__init__`.
+if cfg fails to supply a required argument for
+VanillaMPNNGraphUpdate.__init__.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_get_config_dict.md b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_get_config_dict.md
index 82ee8053..879fba16 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_get_config_dict.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/vanilla_mpnn/graph_update_get_config_dict.md
@@ -1,17 +1,10 @@
 # vanilla_mpnn.graph_update_get_config_dict
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns ConfigDict for graph_update_from_config_dict() with defaults.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/runner.md b/tensorflow_gnn/docs/api_docs/python/runner.md
new file mode 100644
index 00000000..e8a39924
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner.md
@@ -0,0 +1,173 @@
+# Module: runner
+
+
+
+
+ View source
+on GitHub 
+
+A general purpose runner for TF-GNN.
+
+## Classes
+
+[`class ContextLabelFn`](./runner/ContextLabelFn.md): Reads out a `tfgnn.Field`
+from the `GraphTensor` context.
+
+[`class DatasetProvider`](./runner/DatasetProvider.md): Helper class that
+provides a standard way to create an ABC using inheritance.
+
+[`class DotProductLinkPrediction`](./runner/DotProductLinkPrediction.md):
+Implements edge score as dot product of features of endpoint nodes.
+
+[`class FitOrSkipPadding`](./runner/FitOrSkipPadding.md): Calculates fit or skip
+`SizeConstraints` for `GraphTensor` padding.
+
+[`class GraphBinaryClassification`](./runner/GraphBinaryClassification.md):
+Graph binary (or multi-label) classification from pooled node states.
+
+[`class GraphMeanAbsoluteError`](./runner/GraphMeanAbsoluteError.md): Mean
+absolute error task.
+
+[`class GraphMeanAbsolutePercentageError`](./runner/GraphMeanAbsolutePercentageError.md):
+Mean absolute percentage error task.
+
+[`class GraphMeanSquaredError`](./runner/GraphMeanSquaredError.md): Mean squared
+error task.
+
+[`class GraphMeanSquaredLogScaledError`](./runner/GraphMeanSquaredLogScaledError.md):
+Mean squared log scaled error task.
+
+[`class GraphMeanSquaredLogarithmicError`](./runner/GraphMeanSquaredLogarithmicError.md):
+Mean squared logarithmic error task.
+
+[`class GraphMulticlassClassification`](./runner/GraphMulticlassClassification.md):
+Graph multiclass classification from pooled node states.
+
+[`class GraphTensorPadding`](./runner/GraphTensorPadding.md): Collects
+`GraphtTensor` padding helpers.
+
+[`class GraphTensorProcessorFn`](./runner/GraphTensorProcessorFn.md): A class
+for `GraphTensor` processing.
+
+[`class HadamardProductLinkPrediction`](./runner/HadamardProductLinkPrediction.md):
+Implements edge score as hadamard product of features of endpoint nodes.
+
+[`class IntegratedGradientsExporter`](./runner/IntegratedGradientsExporter.md):
+Exports a Keras model with an additional integrated gradients signature.
+
+[`class KerasModelExporter`](./runner/KerasModelExporter.md): Exports a Keras
+model (with Keras API) via `tf.keras.models.save_model`.
+
+[`class KerasTrainer`](./runner/KerasTrainer.md): Trains using the
+`tf.keras.Model.fit` training loop.
+
+[`class KerasTrainerCheckpointOptions`](./runner/KerasTrainerCheckpointOptions.md):
+Provides Keras Checkpointing related configuration options.
+
+[`class KerasTrainerOptions`](./runner/KerasTrainerOptions.md): Provides Keras
+training related options.
+
+[`class ModelExporter`](./runner/ModelExporter.md): Saves a Keras model.
+
+[`class NodeBinaryClassification`](./runner/NodeBinaryClassification.md): Node
+binary (or multi-label) classification via structured readout.
+
+[`class NodeMulticlassClassification`](./runner/NodeMulticlassClassification.md):
+Node multiclass classification via structured readout.
+
+[`class ParameterServerStrategy`](./runner/ParameterServerStrategy.md): A
+`ParameterServerStrategy` convenience wrapper.
+
+[`class PassthruDatasetProvider`](./runner/PassthruDatasetProvider.md): Builds a
+`tf.data.Dataset` from a pass thru dataset.
+
+[`class PassthruSampleDatasetsProvider`](./runner/PassthruSampleDatasetsProvider.md):
+Builds a sampled `tf.data.Dataset` from multiple pass thru datasets.
+
+[`class RootNodeBinaryClassification`](./runner/RootNodeBinaryClassification.md):
+Root node binary (or multi-label) classification.
+
+[`class RootNodeLabelFn`](./runner/RootNodeLabelFn.md): Reads out a
+`tfgnn.Field` from the `GraphTensor` root (i.e. first) node.
+
+[`class RootNodeMeanAbsoluteError`](./runner/RootNodeMeanAbsoluteError.md): Mean
+absolute error task.
+
+[`class RootNodeMeanAbsoluteLogarithmicError`](./runner/RootNodeMeanAbsoluteLogarithmicError.md):
+Root node mean absolute logarithmic error task.
+
+[`class RootNodeMeanAbsolutePercentageError`](./runner/RootNodeMeanAbsolutePercentageError.md):
+Mean absolute percentage error task.
+
+[`class RootNodeMeanSquaredError`](./runner/RootNodeMeanSquaredError.md): Mean
+squared error task.
+
+[`class RootNodeMeanSquaredLogScaledError`](./runner/RootNodeMeanSquaredLogScaledError.md):
+Mean squared log scaled error task.
+
+[`class RootNodeMeanSquaredLogarithmicError`](./runner/RootNodeMeanSquaredLogarithmicError.md):
+Mean squared logarithmic error task.
+
+[`class RootNodeMulticlassClassification`](./runner/RootNodeMulticlassClassification.md):
+Root node multiclass classification.
+
+[`class RunResult`](./runner/RunResult.md): Holds the return values of
+`run(...)`.
+
+[`class SampleTFRecordDatasetsProvider`](./runner/SampleTFRecordDatasetsProvider.md):
+Builds a sampling `tf.data.Dataset` from multiple filenames.
+
+[`class SimpleDatasetProvider`](./runner/SimpleDatasetProvider.md): Builds a
+`tf.data.Dataset` from a list of files.
+
+[`class SimpleSampleDatasetsProvider`](./runner/SimpleSampleDatasetsProvider.md):
+Builds a sampling `tf.data.Dataset` from multiple filenames.
+
+[`class SubmoduleExporter`](./runner/SubmoduleExporter.md): Exports a Keras
+submodule.
+
+[`class TFDataServiceConfig`](./runner/TFDataServiceConfig.md): Provides tf.data
+service related configuration options.
+
+[`class TFRecordDatasetProvider`](./runner/TFRecordDatasetProvider.md): Builds a
+`tf.data.Dataset` from a list of files.
+
+[`class TPUStrategy`](./runner/TPUStrategy.md): A `TPUStrategy` convenience
+wrapper.
+
+[`class Task`](./runner/Task.md): Defines a learning objective for a GNN.
+
+[`class TightPadding`](./runner/TightPadding.md): Calculates tight
+`SizeConstraints` for `GraphTensor` padding.
+
+[`class Trainer`](./runner/Trainer.md): A class for training and validation of a
+Keras model.
+
+## Functions
+
+[`export_model(...)`](./runner/export_model.md): Exports a Keras model without
+traces s.t. it is loadable without TF-GNN.
+
+[`incrementing_model_dir(...)`](./runner/incrementing_model_dir.md): Create,
+given some `dirname`, an incrementing model directory.
+
+[`integrated_gradients(...)`](./runner/integrated_gradients.md): Integrated
+gradients.
+
+[`one_node_per_component(...)`](./runner/one_node_per_component.md): Returns a
+`Mapping` `node_set_name: 1` for every node set in `gtspec`.
+
+[`run(...)`](./runner/run.md): Runs training (and validation) of a model on
+task(s) with the given data.
+
+## Type Aliases
+
+[`Loss`](./runner/Loss.md)
+
+[`Losses`](./runner/Losses.md)
+
+[`Metric`](./runner/Loss.md)
+
+[`Metrics`](./runner/Metrics.md)
+
+[`Predictions`](./runner/Predictions.md)
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/ContextLabelFn.md b/tensorflow_gnn/docs/api_docs/python/runner/ContextLabelFn.md
new file mode 100644
index 00000000..b5145681
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/ContextLabelFn.md
@@ -0,0 +1,17 @@
+# runner.ContextLabelFn
+
+
+
+
+ View source
+on GitHub 
+
+Reads out a `tfgnn.Field` from the `GraphTensor` context.
+
+
+runner.ContextLabelFn(
+    feature_name: str, **kwargs
+)
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/DatasetProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/DatasetProvider.md
new file mode 100644
index 00000000..e7f9b971
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/DatasetProvider.md
@@ -0,0 +1,27 @@
+# runner.DatasetProvider
+
+
+
+
+ View source
+on GitHub 
+
+Helper class that provides a standard way to create an ABC using inheritance.
+
+
+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+@abc.abstractmethod
+get_dataset(
+    context: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Get a `tf.data.Dataset` by `context` per replica.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/DotProductLinkPrediction.md b/tensorflow_gnn/docs/api_docs/python/runner/DotProductLinkPrediction.md
new file mode 100644
index 00000000..dd562c93
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/DotProductLinkPrediction.md
@@ -0,0 +1,200 @@
+# runner.DotProductLinkPrediction
+
+
+
+
+ View source
+on GitHub 
+
+Implements edge score as dot product of features of endpoint nodes.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.DotProductLinkPrediction(
+    *,
+    node_feature_name: tfgnn.FieldName = tfgnn.HIDDEN_STATE,
+    readout_label_feature_name: str = 'label',
+    readout_node_set_name: tfgnn.NodeSetName = '_readout'
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+node_feature_name
+
+Name of feature where node state for link-prediction
+is read from. The final link prediction score will be:
+`score(graph.node_sets[source][node_feature_name],
+        graph.node_sets[target][node_feature_name])`
+where source and target, respectively, are:
+graph.edge_sets[readout_node_set_name+"/source"].adjacency.source_name
+and
+graph.edge_sets[readout_node_set_name+"/target"].adjacency.source_name
+

+readout_label_feature_name
+
+The labels for edge connections,
+source nodes
+graph.edge_sets[readout_node_set_name+"/source"].adjacency.source in
+node set graph.node_sets[source] against target nodes
+graph.edge_sets[readout_node_set_name+"/target"].adjacency.source in
+node set graph.node_sets[source], must be stored in
+graph.node_sets[readout_node_set_name][readout_label_feature_name].
+

+readout_node_set_name
+
+Determines the readout node-set, which must have
+feature readout_label_feature_name, and must receive connections (at
+target endpoints) from edge-sets readout_node_set_name+"/source" and
+readout_node_set_name+"/target".
+

+
+## Methods
+
+
losses

+
+View
+source
+
+
+losses() -> runner.Losses
+

+
+Binary cross-entropy.
+
+
metrics

+
+View
+source
+
+
+metrics() -> runner.Metrics
+

+
+Returns arbitrary task specific metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    graph: tfgnn.GraphTensor
+) -> runner.Predictions
+

+
+Produces prediction outputs for the learning objective.
+
+Overall model composition* makes use of the Keras Functional API
+(https://www.tensorflow.org/guide/keras/functional) to map symbolic Keras
+`GraphTensor` inputs to symbolic Keras `Field` outputs. Outputs must match the
+structure (one or mapping) of labels from `preprocess`.
+
+*) `outputs = predict(GNN(inputs))` where `inputs` are those `GraphTensor`
+returned by `preprocess(...)`, `GNN` is the base GNN, `predict` is this method
+and `outputs` are the prediction outputs for the learning objective.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+*args
+
+The symbolic Keras GraphTensor inputs(s). These inputs correspond
+(in sequence) to the base GNN output of each GraphTensor returned by
+preprocess(...).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The model's prediction output for this task.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    gt: tfgnn.GraphTensor
+) -> Tuple[tfgnn.GraphTensor, tfgnn.Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/FitOrSkipPadding.md b/tensorflow_gnn/docs/api_docs/python/runner/FitOrSkipPadding.md
new file mode 100644
index 00000000..3e349260
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/FitOrSkipPadding.md
@@ -0,0 +1,49 @@
+# runner.FitOrSkipPadding
+
+
+
+
+ View source
+on GitHub 
+
+Calculates fit or skip `SizeConstraints` for `GraphTensor` padding.
+
+Inherits From: [`GraphTensorPadding`](../runner/GraphTensorPadding.md)
+
+
+runner.FitOrSkipPadding(
+    gtspec: tfgnn.GraphTensorSpec,
+    dataset_provider: runner.DatasetProvider,
+    min_nodes_per_component: Optional[Mapping[str, int]] = None,
+    fit_or_skip_sample_sample_size: int = 10000,
+    fit_or_skip_success_ratio: float = 0.99
+)
+

+
+
+
+See: `tfgnn.learn_fit_or_skip_size_constraints.`
+
+## Methods
+
+
get_filter_fn

+
+View
+source
+
+
+get_filter_fn(
+    size_constraints: SizeConstraints
+) -> Callable[..., bool]
+

+
+
get_size_constraints

+
+View
+source
+
+
+get_size_constraints(
+    target_batch_size: int
+) -> SizeConstraints
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphBinaryClassification.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphBinaryClassification.md
new file mode 100644
index 00000000..237ce1c9
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphBinaryClassification.md
@@ -0,0 +1,222 @@
+# runner.GraphBinaryClassification
+
+
+
+
+ View source
+on GitHub 
+
+Graph binary (or multi-label) classification from pooled node states.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphBinaryClassification(
+    node_set_name: str,
+    units: int = 1,
+    *,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    reduce_type: str = 'mean',
+    name: str = 'classification_logits',
+    label_fn: Optional[LabelFn] = None,
+    label_feature_name: Optional[str] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+node_set_name
+
+The node set to pool.
+

+units
+
+The units for the classification head. (Typically 1 for binary
+classification and the number of labels for multi-label classification.)
+

+state_name
+
+The feature name for activations (e.g.: tfgnn.HIDDEN_STATE).
+

+reduce_type
+
+The context pooling reduction type.
+

+name
+
+The classification head's layer name. To control the naming of saved
+model outputs see the runner model exporters (e.g.,
+KerasModelExporter).
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> Field
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+Returns arbitrary task specific losses.
+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Returns arbitrary task specific metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for classification.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for classification.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The classification logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsoluteError.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsoluteError.md
new file mode 100644
index 00000000..0894d92a
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsoluteError.md
@@ -0,0 +1,195 @@
+# runner.GraphMeanAbsoluteError
+
+
+
+
+ View source
+on GitHub 
+
+Mean absolute error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphMeanAbsoluteError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    reduce_type: str = 'mean',
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsolutePercentageError.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsolutePercentageError.md
new file mode 100644
index 00000000..4794f58e
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanAbsolutePercentageError.md
@@ -0,0 +1,195 @@
+# runner.GraphMeanAbsolutePercentageError
+
+
+
+
+ View source
+on GitHub 
+
+Mean absolute percentage error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphMeanAbsolutePercentageError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    reduce_type: str = 'mean',
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredError.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredError.md
new file mode 100644
index 00000000..fc7f9332
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredError.md
@@ -0,0 +1,195 @@
+# runner.GraphMeanSquaredError
+
+
+
+
+ View source
+on GitHub 
+
+Mean squared error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphMeanSquaredError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    reduce_type: str = 'mean',
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogScaledError.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogScaledError.md
new file mode 100644
index 00000000..00e3fe85
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogScaledError.md
@@ -0,0 +1,155 @@
+# runner.GraphMeanSquaredLogScaledError
+
+
+
+
+ View source
+on GitHub 
+
+Mean squared log scaled error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphMeanSquaredLogScaledError(
+    *args,
+    alpha_loss_param: float = 5.0,
+    epsilon_loss_param: float = 1e-08,
+    reduction: tf.keras.losses.Reduction = AUTO,
+    name: Optional[str] = None,
+    **kwargs
+)
+

+
+
+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogarithmicError.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogarithmicError.md
new file mode 100644
index 00000000..3031008c
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphMeanSquaredLogarithmicError.md
@@ -0,0 +1,195 @@
+# runner.GraphMeanSquaredLogarithmicError
+
+
+
+
+ View source
+on GitHub 
+
+Mean squared logarithmic error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphMeanSquaredLogarithmicError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    reduce_type: str = 'mean',
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphMulticlassClassification.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphMulticlassClassification.md
new file mode 100644
index 00000000..e667763c
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphMulticlassClassification.md
@@ -0,0 +1,239 @@
+# runner.GraphMulticlassClassification
+
+
+
+
+ View source
+on GitHub 
+
+Graph multiclass classification from pooled node states.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.GraphMulticlassClassification(
+    node_set_name: str,
+    *,
+    num_classes: Optional[int] = None,
+    class_names: Optional[Sequence[str]] = None,
+    per_class_statistics: bool = False,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    reduce_type: str = 'mean',
+    name: str = 'classification_logits',
+    label_fn: Optional[LabelFn] = None,
+    label_feature_name: Optional[str] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+node_set_name
+
+The node set to pool.
+

+num_classes
+
+The number of classes. Exactly one of num_classes or
+class_names must be specified
+

+class_names
+
+The class names. Exactly one of num_classes or
+class_names must be specified
+

+per_class_statistics
+
+Whether to compute statistics per class.
+

+state_name
+
+The feature name for activations (e.g.: tfgnn.HIDDEN_STATE).
+

+reduce_type
+
+The context pooling reduction type.
+

+name
+
+The classification head's layer name. To control the naming of saved
+model outputs see the runner model exporters (e.g.,
+KerasModelExporter).
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> Field
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+Sparse categorical crossentropy loss.
+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Sparse categorical metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for classification.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for classification.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The classification logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphTensorPadding.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphTensorPadding.md
new file mode 100644
index 00000000..63ece9ef
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphTensorPadding.md
@@ -0,0 +1,37 @@
+# runner.GraphTensorPadding
+
+
+
+
+ View source
+on GitHub 
+
+Collects `GraphtTensor` padding helpers.
+
+
+
+## Methods
+
+
get_filter_fn

+
+View
+source
+
+
+@abc.abstractmethod
+get_filter_fn(
+    size_constraints: SizeConstraints
+) -> Callable[..., bool]
+

+
+
get_size_constraints

+
+View
+source
+
+
+@abc.abstractmethod
+get_size_constraints(
+    target_batch_size: int
+) -> SizeConstraints
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/GraphTensorProcessorFn.md b/tensorflow_gnn/docs/api_docs/python/runner/GraphTensorProcessorFn.md
new file mode 100644
index 00000000..55b2eb0e
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/GraphTensorProcessorFn.md
@@ -0,0 +1,26 @@
+# runner.GraphTensorProcessorFn
+
+
+
+
+ View source
+on GitHub 
+
+A class for `GraphTensor` processing.
+
+
+
+## Methods
+
+
__call__

+
+View
+source
+
+
+__call__(
+    inputs: GraphTensor
+) -> GraphTensor
+

+
+Processes a `GraphTensor`.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/HadamardProductLinkPrediction.md b/tensorflow_gnn/docs/api_docs/python/runner/HadamardProductLinkPrediction.md
new file mode 100644
index 00000000..d66469ba
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/HadamardProductLinkPrediction.md
@@ -0,0 +1,203 @@
+# runner.HadamardProductLinkPrediction
+
+
+
+
+ View source
+on GitHub 
+
+Implements edge score as hadamard product of features of endpoint nodes.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.HadamardProductLinkPrediction(
+    *,
+    node_feature_name: tfgnn.FieldName = tfgnn.HIDDEN_STATE,
+    readout_label_feature_name: str = 'label',
+    readout_node_set_name: tfgnn.NodeSetName = '_readout'
+)
+

+
+
+
+The hadamard product is followed by one layer with scalar output.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+node_feature_name
+
+Name of feature where node state for link-prediction
+is read from. The final link prediction score will be:
+`score(graph.node_sets[source][node_feature_name],
+        graph.node_sets[target][node_feature_name])`
+where source and target, respectively, are:
+graph.edge_sets[readout_node_set_name+"/source"].adjacency.source_name
+and
+graph.edge_sets[readout_node_set_name+"/target"].adjacency.source_name
+

+readout_label_feature_name
+
+The labels for edge connections,
+source nodes
+graph.edge_sets[readout_node_set_name+"/source"].adjacency.source in
+node set graph.node_sets[source] against target nodes
+graph.edge_sets[readout_node_set_name+"/target"].adjacency.source in
+node set graph.node_sets[source], must be stored in
+graph.node_sets[readout_node_set_name][readout_label_feature_name].
+

+readout_node_set_name
+
+Determines the readout node-set, which must have
+feature readout_label_feature_name, and must receive connections (at
+target endpoints) from edge-sets readout_node_set_name+"/source" and
+readout_node_set_name+"/target".
+

+
+## Methods
+
+
losses

+
+View
+source
+
+
+losses() -> runner.Losses
+

+
+Binary cross-entropy.
+
+
metrics

+
+View
+source
+
+
+metrics() -> runner.Metrics
+

+
+Returns arbitrary task specific metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    graph: tfgnn.GraphTensor
+) -> runner.Predictions
+

+
+Produces prediction outputs for the learning objective.
+
+Overall model composition* makes use of the Keras Functional API
+(https://www.tensorflow.org/guide/keras/functional) to map symbolic Keras
+`GraphTensor` inputs to symbolic Keras `Field` outputs. Outputs must match the
+structure (one or mapping) of labels from `preprocess`.
+
+*) `outputs = predict(GNN(inputs))` where `inputs` are those `GraphTensor`
+returned by `preprocess(...)`, `GNN` is the base GNN, `predict` is this method
+and `outputs` are the prediction outputs for the learning objective.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+*args
+
+The symbolic Keras GraphTensor inputs(s). These inputs correspond
+(in sequence) to the base GNN output of each GraphTensor returned by
+preprocess(...).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The model's prediction output for this task.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    gt: tfgnn.GraphTensor
+) -> Tuple[tfgnn.GraphTensor, tfgnn.Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/IntegratedGradientsExporter.md b/tensorflow_gnn/docs/api_docs/python/runner/IntegratedGradientsExporter.md
new file mode 100644
index 00000000..2aac8e65
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/IntegratedGradientsExporter.md
@@ -0,0 +1,127 @@
+# runner.IntegratedGradientsExporter
+
+
+
+
+ View source
+on GitHub 
+
+Exports a Keras model with an additional integrated gradients signature.
+
+Inherits From: [`ModelExporter`](../runner/ModelExporter.md)
+
+
+runner.IntegratedGradientsExporter(
+    integrated_gradients_output_name: Optional[str] = None,
+    subdirectory: Optional[str] = None,
+    random_counterfactual: bool = True,
+    steps: int = 32,
+    seed: Optional[int] = None,
+    options: Optional[tf.saved_model.SaveOptions] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+integrated_gradients_output_name
+
+The name for the integrated gradients
+output tensor. If unset, the tensor will be named by Keras defaults.
+

+subdirectory
+
+An optional subdirectory, if set: models are exported to
+os.path.join(export_dir, subdirectory).
+

+random_counterfactual
+
+Whether to use a random uniform counterfactual.
+

+steps
+
+The number of interpolations of the Riemann sum approximation.
+

+seed
+
+An optional random seed.
+

+options
+
+Options for saving to SavedModel.
+

+
+## Methods
+
+
save

+
+View
+source
+
+
+save(
+    run_result: runner.RunResult,
+    export_dir: str
+)
+

+
+Exports a Keras model with an additional integrated gradients signature.
+
+Importantly: the `run_result.preprocess_model`, if provided, and
+`run_result.trained_model` are stacked before any export. Stacking involves the
+chaining of the first output of `run_result.preprocess_model` to the only input
+of `run_result.trained_model.` The result is a model with the input of
+`run_result.preprocess_model` and the output of `run_result.trained_model.`
+
+Two serving signatures are exported:
+
+'serving_default') The default serving signature (i.e., the `preprocess_model`
+input signature), 'integrated_gradients') The integrated gradients signature
+(i.e., the `preprocess_model` input signature).
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+run_result
+
+A RunResult from training.
+

+export_dir
+
+A destination directory.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/KerasModelExporter.md b/tensorflow_gnn/docs/api_docs/python/runner/KerasModelExporter.md
new file mode 100644
index 00000000..4de9ca64
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/KerasModelExporter.md
@@ -0,0 +1,121 @@
+# runner.KerasModelExporter
+
+
+
+
+ View source
+on GitHub 
+
+Exports a Keras model (with Keras API) via `tf.keras.models.save_model`.
+
+Inherits From: [`ModelExporter`](../runner/ModelExporter.md)
+
+
+runner.KerasModelExporter(
+    *,
+    output_names: Optional[Any] = None,
+    subdirectory: Optional[str] = None,
+    include_preprocessing: bool = True,
+    options: Optional[tf.saved_model.SaveOptions] = None,
+    use_legacy_model_save: Optional[bool] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+output_names
+
+By default, each output of the exported model uses the name
+of the final Keras layer that created it as its key in the SavedModel
+signature. This argument can be set to a single str name or a nested
+structure of str names to override the output names. Its nesting
+structure must match the exported model's output (as checked by
+tf.nest.assert_same_structure). Any None values in output_names
+are ignored, leaving that output with its default name.
+

+subdirectory
+
+An optional subdirectory, if set: models are exported to
+os.path.join(export_dir, subdirectory).
+

+include_preprocessing
+
+Whether to include any preprocess_model.
+

+options
+
+Options for saving to a TensorFlow SavedModel.
+

+use_legacy_model_save
+
+Optional; most users can leave it unset to get a
+useful default for export to inference. See runner.export_model()
+for more.
+

+
+## Methods
+
+
save

+
+View
+source
+
+
+save(
+    run_result: runner.RunResult,
+    export_dir: str
+)
+

+
+Exports a Keras model (with Keras API) via tf.keras.models.save_model.
+
+Importantly: the `run_result.preprocess_model`, if provided, and
+`run_result.trained_model` are stacked before any export. Stacking involves the
+chaining of the first output of `run_result.preprocess_model` to the only input
+of `run_result.trained_model.` The result is a model with the input of
+`run_result.preprocess_model` and the output of `run_result.trained_model.`
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+run_result
+
+A RunResult from training.
+

+export_dir
+
+A destination directory.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainer.md b/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainer.md
new file mode 100644
index 00000000..a8468368
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainer.md
@@ -0,0 +1,263 @@
+# runner.KerasTrainer
+
+
+
+
+ View source
+on GitHub 
+
+Trains using the `tf.keras.Model.fit` training loop.
+
+Inherits From: [`Trainer`](../runner/Trainer.md)
+
+
+runner.KerasTrainer(
+    strategy: tf.distribute.Strategy,
+    *,
+    model_dir: str,
+    checkpoint_options: Optional[runner.KerasTrainerCheckpointOptions] = None,
+    backup_dir: Optional[str] = None,
+    steps_per_epoch: Optional[int] = None,
+    verbose: Union[int, str] = 'auto',
+    validation_steps: Optional[int] = None,
+    validation_per_epoch: Optional[int] = None,
+    validation_freq: Optional[int] = None,
+    summarize_every_n_steps: Union[int, str] = 500,
+    checkpoint_every_n_steps: Union[int, str] = 'epoch',
+    backup_and_restore: bool = True,
+    callbacks: Optional[Sequence[tf.keras.callbacks.Callback]] = None,
+    restore_best_weights: Optional[bool] = None,
+    options: Optional[runner.KerasTrainerOptions] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+strategy
+
+A tf.distribute.Strategy.
+

+model_dir
+
+A model directory for summaries.
+

+checkpoint_options
+
+An optional configuration for checkpointing related
+configs. If checkpoint_options.checkpoint_dir is unset;
+os.path.join(model_dir, "ckpnt") is used.
+

+backup_dir
+
+An optional directory for backup, if unset;
+(os.path.join(model_dir, "backup"),) is used.
+

+steps_per_epoch
+
+The number of training steps per epoch. Optional,
+if unspecified: epochs are at tf.data.Dataset end.
+

+verbose
+
+Forwarded to tf.keras.Model.fit(). Possible values are
+0 (silent), 1 (print progress bar), 2 (one line per epoch), and
+"auto" (default) defers to keras to select verbosity.
+

+validation_steps
+
+The number of steps used during validation. Optional,
+if unspecified: the entire validation tf.data.Dataset is evaluated.
+

+validation_per_epoch
+
+The number of validations done per training epoch.
+Optional, if unspecified: Perform one validation per training epoch.
+Only one of validation_per_epoch and validation_freq can be
+specified.
+

+validation_freq
+
+Specifies how many training epochs to run before a new
+validation run is performed. Optional, if unspecified: Performs
+validation after every training epoch. Only one of
+validation_per_epoch and validation_freq can be specified.
+

+summarize_every_n_steps
+
+The frequency for writing TensorBoard summaries,
+as an integer number of steps, or "epoch" for once per epoch, or
+"never".
+

+checkpoint_every_n_steps
+
+The frequency for writing latest models, as an
+integer number of steps, or "epoch" for once per epoch, or "never".
+The best model will always be saved after each validation epoch except
+when this parameter is set to "never", because the validation metric is
+available only after validation epoch.
+

+backup_and_restore
+
+Whether to backup and restore (According to
+tf.keras.callbacks.BackupAndRestore). The backup
+directory is determined by backup_dir.
+

+callbacks
+
+Optional additional tf.keras.callbacks.Callback for
+tf.keras.Model.fit.
+

+restore_best_weights
+
+Requires a checkpoint_every_n_steps other than
+"never." Whether to restore the best model weights as determined by
+tf.keras.callbacks.ModelCheckpoint after training. If unspecified,
+its value is determined at train(...) invocation: True if
+valid_ds_provider is not None else False.
+

+options
+
+A KerasTrainerOptions.
+

+
+
+
+ 
+
+
+
+
+
+
Attributes
  model_dir  
+
+ 
  strategy  
+
+

+
+## Methods
+
+
train

+
+View
+source
+
+
+train(
+    model_fn: Callable[[], tf.keras.Model],
+    train_ds_provider: runner.DatasetProvider,
+    *,
+    epochs: int = 1,
+    valid_ds_provider: Optional[runner.DatasetProvider] = None
+) -> tf.keras.Model
+

+
+Runs `tf.keras.Model.fit` with the`tf.distribute.Strategy` provided.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+model_fn
+
+A ModelFn, to be invoked in the tf.distribute.Strategty
+scope.
+

+train_ds_provider
+
+A function that returns a tf.data.Dataset for
+training.The items of the tf.data.Dataset are pairs
+(graph_tensor, label) that represent one batch of per-replica training
+inputs after GraphTensor.merge_batch_to_components() has been applied.
+

+epochs
+
+The epochs to train: adjusted for validation_per_epoch.
+

+valid_ds_provider
+
+An optional function that returns a tf.data.Dataset
+for validation. The items of the tf.data.Dataset are pairs
+(graph_tensor, label) that represent one batch of per-replica training
+inputs after GraphTensor.merge_batch_to_components() has been applied.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A trained tf.keras.Model.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerCheckpointOptions.md b/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerCheckpointOptions.md
new file mode 100644
index 00000000..b3fbff8d
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerCheckpointOptions.md
@@ -0,0 +1,108 @@
+# runner.KerasTrainerCheckpointOptions
+
+
+
+
+ View source
+on GitHub 
+
+Provides Keras Checkpointing related configuration options.
+
+
+runner.KerasTrainerCheckpointOptions(
+    checkpoint_dir: Optional[str] = None,
+    best_checkpoint: str = 'best',
+    latest_checkpoint: str = 'latest'
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+checkpoint_dir
+
+Directory path to save checkpoint files.
+

+best_checkpoint
+
+Filename for the best checkpoint.
+

+latest_checkpoint
+
+Filename for the latest checkpoint.
+

+
+## Methods
+
+
best_checkpoint_filepath

+
+View
+source
+
+
+best_checkpoint_filepath() -> str
+

+
+
latest_checkpoint_filepath

+
+View
+source
+
+
+latest_checkpoint_filepath() -> str
+

+
+
__eq__

+
+
+__eq__(
+    other
+)
+

+
+Return self==value.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Class Variables

+best_checkpoint
+
+'best'
+

+checkpoint_dir
+
+None
+

+latest_checkpoint
+
+'latest'
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerOptions.md b/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerOptions.md
new file mode 100644
index 00000000..fd30bc44
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/KerasTrainerOptions.md
@@ -0,0 +1,90 @@
+# runner.KerasTrainerOptions
+
+
+
+
+ View source
+on GitHub 
+
+Provides Keras training related options.
+
+
+runner.KerasTrainerOptions(
+    policy: Optional[Union[str, tf.keras.mixed_precision.Policy]] = None,
+    soft_device_placement: bool = False,
+    enable_check_numerics: bool = False
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+policy
+
+Dataclass field
+

+soft_device_placement
+
+Dataclass field
+

+enable_check_numerics
+
+Dataclass field
+

+
+## Methods
+
+
__eq__

+
+
+__eq__(
+    other
+)
+

+
+Return self==value.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Class Variables

+enable_check_numerics
+
+False
+

+policy
+
+None
+

+soft_device_placement
+
+False
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/Loss.md b/tensorflow_gnn/docs/api_docs/python/runner/Loss.md
new file mode 100644
index 00000000..fa45eb7e
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/Loss.md
@@ -0,0 +1,17 @@
+# runner.Loss
+
+
+
+This symbol is a **type alias**.
+
+#### Source:
+
+
+Loss = Callable[
+    tf.Tensor,
+    tf.Tensor,
+    tf.Tensor
+]
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/Losses.md b/tensorflow_gnn/docs/api_docs/python/runner/Losses.md
new file mode 100644
index 00000000..839e40c3
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/Losses.md
@@ -0,0 +1,16 @@
+# runner.Losses
+
+
+
+This symbol is a **type alias**.
+
+#### Source:
+
+
+Losses = Union[
+    runner.Loss,
+    Mapping[str, runner.Loss]
+]
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/Metrics.md b/tensorflow_gnn/docs/api_docs/python/runner/Metrics.md
new file mode 100644
index 00000000..5d281a66
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/Metrics.md
@@ -0,0 +1,17 @@
+# runner.Metrics
+
+
+
+This symbol is a **type alias**.
+
+#### Source:
+
+
+Metrics = Union[
+    runner.Loss,
+    Sequence[runner.Loss],
+    Mapping[str, Union[runner.Loss, Sequence[runner.Loss]]]
+]
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/ModelExporter.md b/tensorflow_gnn/docs/api_docs/python/runner/ModelExporter.md
new file mode 100644
index 00000000..3fbc5b4c
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/ModelExporter.md
@@ -0,0 +1,53 @@
+# runner.ModelExporter
+
+
+
+
+ View source
+on GitHub 
+
+Saves a Keras model.
+
+
+
+## Methods
+
+
save

+
+View
+source
+
+
+@abc.abstractmethod
+save(
+    run_result: RunResult, export_dir: str
+)
+

+
+Saves a Keras model.
+
+All persistence decisions are left to the implementation: e.g., a Keras model
+with full API or a simple `tf.train.Checkpoint` may be saved.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+run_result
+
+A RunResult from training.
+

+export_dir
+
+A destination directory.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/NodeBinaryClassification.md b/tensorflow_gnn/docs/api_docs/python/runner/NodeBinaryClassification.md
new file mode 100644
index 00000000..2926e68c
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/NodeBinaryClassification.md
@@ -0,0 +1,237 @@
+# runner.NodeBinaryClassification
+
+
+
+
+ View source
+on GitHub 
+
+Node binary (or multi-label) classification via structured readout.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.NodeBinaryClassification(
+    key: str = 'seed',
+    units: int = 1,
+    *,
+    feature_name: str = tfgnn.HIDDEN_STATE,
+    readout_node_set: tfgnn.NodeSetName = '_readout',
+    validate: bool = True,
+    name: str = 'classification_logits',
+    label_fn: Optional[LabelFn] = None,
+    label_feature_name: Optional[str] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+key
+
+A string key to select between possibly multiple named readouts.
+

+units
+
+The units for the classification head. (Typically 1 for binary
+classification and the number of labels for multi-label classification.)
+

+feature_name
+
+The name of the feature to read. If unset,
+tfgnn.HIDDEN_STATE will be read.
+

+readout_node_set
+
+A string, defaults to "_readout". This is used as the
+name for the readout node set and as a name prefix for its edge sets.
+

+validate
+
+Setting this to false disables the validity checks for the
+auxiliary edge sets. This is stronlgy discouraged, unless great care is
+taken to run tfgnn.validate_graph_tensor_for_readout() earlier on
+structurally unchanged GraphTensors.
+

+name
+
+The classification head's layer name. To control the naming of saved
+model outputs see the runner model exporters (e.g.,
+KerasModelExporter).
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> Field
+

+
+Gather activations from auxiliary node (and edge) sets.
+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+Returns arbitrary task specific losses.
+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Returns arbitrary task specific metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for classification.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for classification.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The classification logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/NodeMulticlassClassification.md b/tensorflow_gnn/docs/api_docs/python/runner/NodeMulticlassClassification.md
new file mode 100644
index 00000000..308eea6c
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/NodeMulticlassClassification.md
@@ -0,0 +1,254 @@
+# runner.NodeMulticlassClassification
+
+
+
+
+ View source
+on GitHub 
+
+Node multiclass classification via structured readout.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.NodeMulticlassClassification(
+    key: str = 'seed',
+    *,
+    feature_name: str = tfgnn.HIDDEN_STATE,
+    readout_node_set: tfgnn.NodeSetName = '_readout',
+    validate: bool = True,
+    num_classes: Optional[int] = None,
+    class_names: Optional[Sequence[str]] = None,
+    per_class_statistics: bool = False,
+    name: str = 'classification_logits',
+    label_fn: Optional[LabelFn] = None,
+    label_feature_name: Optional[str] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+key
+
+A string key to select between possibly multiple named readouts.
+

+feature_name
+
+The name of the feature to read. If unset,
+tfgnn.HIDDEN_STATE will be read.
+

+readout_node_set
+
+A string, defaults to "_readout". This is used as the
+name for the readout node set and as a name prefix for its edge sets.
+

+validate
+
+Setting this to false disables the validity checks for the
+auxiliary edge sets. This is stronlgy discouraged, unless great care is
+taken to run tfgnn.validate_graph_tensor_for_readout() earlier on
+structurally unchanged GraphTensors.
+

+num_classes
+
+The number of classes. Exactly one of num_classes or
+class_names must be specified
+

+class_names
+
+The class names. Exactly one of num_classes or
+class_names must be specified
+

+per_class_statistics
+
+Whether to compute statistics per class.
+

+name
+
+The classification head's layer name. To control the naming of saved
+model outputs see the runner model exporters (e.g.,
+KerasModelExporter).
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> Field
+

+
+Gather activations from auxiliary node (and edge) sets.
+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+Sparse categorical crossentropy loss.
+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Sparse categorical metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for classification.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for classification.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The classification logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/ParameterServerStrategy.md b/tensorflow_gnn/docs/api_docs/python/runner/ParameterServerStrategy.md
new file mode 100644
index 00000000..342da409
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/ParameterServerStrategy.md
@@ -0,0 +1,1082 @@
+# runner.ParameterServerStrategy
+
+
+
+
+ View source
+on GitHub 
+
+A `ParameterServerStrategy` convenience wrapper.
+
+
+runner.ParameterServerStrategy(
+    min_shard_bytes: Optional[int] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+    
+    
Args
  cluster_resolver  
+a tf.distribute.cluster_resolver.ClusterResolver object. 
 
+variable_partitioner   a
+distribute.experimental.partitioners.Partitioner that specifies how
+to partition variables. If None, variables will not be partitioned.
+
+*   Predefined partitioners in
+    tf.distribute.experimental.partitioners can be used for this
+    argument. A commonly used partitioner is
+    MinSizePartitioner(min_shard_bytes = 256 << 10, max_shards =
+    num_ps), which allocates at least 256K per shard, and each ps gets at
+    most one shard.
+
+*   variable_partitioner will be called for each variable created
+    under strategy scope to instruct how the variable should be
+    partitioned. Variables that have only one partition along the partitioning
+    axis (i.e., no need for partition) will be created as a normal
+    tf.Variable.
+
+*   Only the first / outermost axis partitioning is supported.
+
+*   Div partition strategy is used to partition variables. Assuming we assign
+    consecutive integer ids along the first axis of a variable, then ids are
+    assigned to shards in a contiguous manner, while attempting to keep each
+    shard size identical. If the ids do not evenly divide the number of shards,
+    each of the first several shards will be assigned one more id. For instance,
+    a variable whose first dimension is 13 has 13 ids, and they are split across
+    5 shards as: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11,
+    12]].
+
+*   Variables created under strategy.extended.colocate_vars_with will
+    not be partitioned.
+    

+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
Attributes
  cluster_resolver  
+Returns the cluster resolver associated with this strategy.
+
+In general, when using a multi-worker tf.distribute strategy such
+as tf.distribute.experimental.MultiWorkerMirroredStrategy or
+tf.distribute.TPUStrategy(), there is a
+tf.distribute.cluster_resolver.ClusterResolver associated with the
+strategy used, and such an instance is returned by this property.
+
+Strategies that intend to have an associated
+tf.distribute.cluster_resolver.ClusterResolver must set the
+relevant attribute, or override this property; otherwise, None is
+returned by default. Those strategies should also provide information regarding
+what is returned by this property.
+
+Single-worker strategies usually do not have a
+tf.distribute.cluster_resolver.ClusterResolver, and in those cases
+this property will return None.
+
+The tf.distribute.cluster_resolver.ClusterResolver may be useful
+when the user needs to access information such as the cluster spec, task type or
+task id. For example,
+
+```python
+
+os.environ['TF_CONFIG'] = json.dumps({
+  'cluster': {
+      'worker': ["localhost:12345", "localhost:23456"],
+      'ps': ["localhost:34567"]
+  },
+  'task': {'type': 'worker', 'index': 0}
+})
+
+# This implicitly uses TF_CONFIG for the cluster and current task info.
+strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+
+...
+
+if strategy.cluster_resolver.task_type == 'worker':
+  # Perform something that's only applicable on workers. Since we set this
+  # as a worker above, this block will run on this particular instance.
+elif strategy.cluster_resolver.task_type == 'ps':
+  # Perform something that's only applicable on parameter servers. Since we
+  # set this as a worker above, this block will not run on this particular
+  # instance.
+```
+
+For more information, please see
+tf.distribute.cluster_resolver.ClusterResolver's API docstring.
+

+extended
+
+tf.distribute.StrategyExtended with additional methods.
+

+num_replicas_in_sync
+
+Returns number of replicas over which gradients are aggregated.
+

+
+## Methods
+
+
distribute_datasets_from_function

+
+
+distribute_datasets_from_function(
+    dataset_fn, options=None
+)
+

+
+Distributes `tf.data.Dataset` instances created by calls to `dataset_fn`.
+
+The argument `dataset_fn` that users pass in is an input function that has a
+`tf.distribute.InputContext` argument and returns a `tf.data.Dataset` instance.
+It is expected that the returned dataset from `dataset_fn` is already batched by
+per-replica batch size (i.e. global batch size divided by the number of replicas
+in sync) and sharded. `tf.distribute.Strategy.distribute_datasets_from_function`
+does not batch or shard the `tf.data.Dataset` instance returned from the input
+function. `dataset_fn` will be called on the CPU device of each of the workers
+and each generates a dataset where every replica on that worker will dequeue one
+batch of inputs (i.e. if a worker has two replicas, two batches will be dequeued
+from the `Dataset` every step).
+
+This method can be used for several purposes. First, it allows you to specify
+your own batching and sharding logic. (In contrast,
+`tf.distribute.experimental_distribute_dataset` does batching and sharding for
+you.) For example, where `experimental_distribute_dataset` is unable to shard
+the input files, this method might be used to manually shard the dataset
+(avoiding the slow fallback behavior in `experimental_distribute_dataset`). In
+cases where the dataset is infinite, this sharding can be done by creating
+dataset replicas that differ only in their random seed.
+
+The `dataset_fn` should take an `tf.distribute.InputContext` instance where
+information about batching and input replication can be accessed.
+
+You can use `element_spec` property of the `tf.distribute.DistributedDataset`
+returned by this API to query the `tf.TypeSpec` of the elements returned by the
+iterator. This can be used to set the `input_signature` property of a
+`tf.function`. Follow `tf.distribute.DistributedDataset.element_spec` to see an
+example.
+
+IMPORTANT: The `tf.data.Dataset` returned by `dataset_fn` should have a
+per-replica batch size, unlike `experimental_distribute_dataset`, which uses the
+global batch size. This may be computed using
+`input_context.get_per_replica_batch_size`.
+
+Note: If you are using TPUStrategy, the order in which the data is processed by
+the workers when using `tf.distribute.Strategy.experimental_distribute_dataset`
+or `tf.distribute.Strategy.distribute_datasets_from_function` is not guaranteed.
+This is typically required if you are using `tf.distribute` to scale prediction.
+You can however insert an index for each element in the batch and order outputs
+accordingly. Refer to
+[this snippet](https://www.tensorflow.org/tutorials/distribute/input#caveats)
+for an example of how to order outputs.
+
+Note: Stateful dataset transformations are currently not supported with
+`tf.distribute.experimental_distribute_dataset` or
+`tf.distribute.distribute_datasets_from_function`. Any stateful ops that the
+dataset may have are currently ignored. For example, if your dataset has a
+`map_fn` that uses `tf.random.uniform` to rotate an image, then you have a
+dataset graph that depends on state (i.e the random seed) on the local machine
+where the python process is being executed.
+
+For a tutorial on more usage and properties of this method, refer to the
+[tutorial on distributed input](https://www.tensorflow.org/tutorials/distribute/input#tfdistributestrategyexperimental_distribute_datasets_from_function)).
+If you are interested in last partial batch handling, read
+[this section](https://www.tensorflow.org/tutorials/distribute/input#partial_batches).
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+dataset_fn
+
+A function taking a tf.distribute.InputContext instance and
+returning a tf.data.Dataset.
+

+options
+
+tf.distribute.InputOptions used to control options on how this
+dataset is distributed.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.distribute.DistributedDataset.
+

+
+
experimental_distribute_dataset

+
+
+experimental_distribute_dataset(
+    dataset, options=None
+)
+

+
+Creates `tf.distribute.DistributedDataset` from `tf.data.Dataset`.
+
+The returned `tf.distribute.DistributedDataset` can be iterated over similar to
+regular datasets. NOTE: The user cannot add any more transformations to a
+`tf.distribute.DistributedDataset`. You can only create an iterator or examine
+the `tf.TypeSpec` of the data generated by it. See API docs of
+`tf.distribute.DistributedDataset` to learn more.
+
+The following is an example:
+
+```
+>>> global_batch_size = 2
+>>> # Passing the devices is optional.
+... strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
+>>> # Create a dataset
+... dataset = tf.data.Dataset.range(4).batch(global_batch_size)
+>>> # Distribute that dataset
+... dist_dataset = strategy.experimental_distribute_dataset(dataset)
+>>> @tf.function
+... def replica_fn(input):
+...   return input*2
+>>> result = []
+>>> # Iterate over the `tf.distribute.DistributedDataset`
+... for x in dist_dataset:
+...   # process dataset elements
+...   result.append(strategy.run(replica_fn, args=(x,)))
+>>> print(result)
+[PerReplica:{
+  0: ,
+  1: 
+}, PerReplica:{
+  0: ,
+  1: 
+}]
+```
+
+Three key actions happening under the hood of this method are batching,
+sharding, and prefetching.
+
+In the code snippet above, `dataset` is batched by `global_batch_size`, and
+calling `experimental_distribute_dataset` on it rebatches `dataset` to a new
+batch size that is equal to the global batch size divided by the number of
+replicas in sync. We iterate through it using a Pythonic for loop. `x` is a
+`tf.distribute.DistributedValues` containing data for all replicas, and each
+replica gets data of the new batch size. `tf.distribute.Strategy.run` will take
+care of feeding the right per-replica data in `x` to the right `replica_fn`
+executed on each replica.
+
+Sharding contains autosharding across multiple workers and within every worker.
+First, in multi-worker distributed training (i.e. when you use
+`tf.distribute.experimental.MultiWorkerMirroredStrategy` or
+`tf.distribute.TPUStrategy`), autosharding a dataset over a set of workers means
+that each worker is assigned a subset of the entire dataset (if the right
+`tf.data.experimental.AutoShardPolicy` is set). This is to ensure that at each
+step, a global batch size of non-overlapping dataset elements will be processed
+by each worker. Autosharding has a couple of different options that can be
+specified using `tf.data.experimental.DistributeOptions`. Then, sharding within
+each worker means the method will split the data among all the worker devices
+(if more than one a present). This will happen regardless of multi-worker
+autosharding.
+
+Note: for autosharding across multiple workers, the default mode is
+`tf.data.experimental.AutoShardPolicy.AUTO`. This mode will attempt to shard the
+input dataset by files if the dataset is being created out of reader datasets
+(e.g. `tf.data.TFRecordDataset`, `tf.data.TextLineDataset`, etc.) or otherwise
+shard the dataset by data, where each of the workers will read the entire
+dataset and only process the shard assigned to it. However, if you have less
+than one input file per worker, we suggest that you disable dataset autosharding
+across workers by setting the
+`tf.data.experimental.DistributeOptions.auto_shard_policy` to be
+`tf.data.experimental.AutoShardPolicy.OFF`.
+
+By default, this method adds a prefetch transformation at the end of the user
+provided `tf.data.Dataset` instance. The argument to the prefetch transformation
+which is `buffer_size` is equal to the number of replicas in sync.
+
+If the above batch splitting and dataset sharding logic is undesirable, please
+use `tf.distribute.Strategy.distribute_datasets_from_function` instead, which
+does not do any automatic batching or sharding for you.
+
+Note: If you are using TPUStrategy, the order in which the data is processed by
+the workers when using `tf.distribute.Strategy.experimental_distribute_dataset`
+or `tf.distribute.Strategy.distribute_datasets_from_function` is not guaranteed.
+This is typically required if you are using `tf.distribute` to scale prediction.
+You can however insert an index for each element in the batch and order outputs
+accordingly. Refer to
+[this snippet](https://www.tensorflow.org/tutorials/distribute/input#caveats)
+for an example of how to order outputs.
+
+Note: Stateful dataset transformations are currently not supported with
+`tf.distribute.experimental_distribute_dataset` or
+`tf.distribute.distribute_datasets_from_function`. Any stateful ops that the
+dataset may have are currently ignored. For example, if your dataset has a
+`map_fn` that uses `tf.random.uniform` to rotate an image, then you have a
+dataset graph that depends on state (i.e the random seed) on the local machine
+where the python process is being executed.
+
+For a tutorial on more usage and properties of this method, refer to the
+[tutorial on distributed input](https://www.tensorflow.org/tutorials/distribute/input#tfdistributestrategyexperimental_distribute_dataset).
+If you are interested in last partial batch handling, read
+[this section](https://www.tensorflow.org/tutorials/distribute/input#partial_batches).
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+dataset
+
+tf.data.Dataset that will be sharded across all replicas using
+the rules stated above.
+

+options
+
+tf.distribute.InputOptions used to control options on how this
+dataset is distributed.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.distribute.DistributedDataset.
+

+
+
experimental_distribute_values_from_function

+
+
+experimental_distribute_values_from_function(
+    value_fn
+)
+

+
+Generates `tf.distribute.DistributedValues` from `value_fn`.
+
+This function is to generate `tf.distribute.DistributedValues` to pass into
+`run`, `reduce`, or other methods that take distributed values when not using
+datasets.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+value_fn
+
+The function to run to generate values. It is called for
+each replica with tf.distribute.ValueContext as the sole argument. It
+must return a Tensor or a type that can be converted to a Tensor.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.distribute.DistributedValues containing a value for each replica.
+

+
+#### Example usage:
+
+1.  Return constant value per replica:
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> def value_fn(ctx):
+    ...   return tf.constant(1.)
+    >>> distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...        value_fn))
+    >>> local_result = strategy.experimental_local_results(
+    ...     distributed_values)
+    >>> local_result
+    (,
+    )
+    ```
+
+2.  Distribute values in array based on replica_id: {: value=2}
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> array_value = np.array([3., 2., 1.])
+    >>> def value_fn(ctx):
+    ...   return array_value[ctx.replica_id_in_sync_group]
+    >>> distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...         value_fn))
+    >>> local_result = strategy.experimental_local_results(
+    ...     distributed_values)
+    >>> local_result
+    (3.0, 2.0)
+    ```
+
+3.  Specify values using num_replicas_in_sync: {: value=3}
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> def value_fn(ctx):
+    ...   return ctx.num_replicas_in_sync
+    >>> distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...         value_fn))
+    >>> local_result = strategy.experimental_local_results(
+    ...     distributed_values)
+    >>> local_result
+    (2, 2)
+    ```
+
+4.  Place values on devices and distribute: {: value=4}
+
+    ```
+    strategy = tf.distribute.TPUStrategy()
+    worker_devices = strategy.extended.worker_devices
+    multiple_values = []
+    for i in range(strategy.num_replicas_in_sync):
+      with tf.device(worker_devices[i]):
+        multiple_values.append(tf.constant(1.0))
+
+    def value_fn(ctx):
+      return multiple_values[ctx.replica_id_in_sync_group]
+
+    distributed_values = strategy.
+      experimental_distribute_values_from_function(
+      value_fn)
+    ```
+
+
experimental_local_results

+
+
+experimental_local_results(
+    value
+)
+

+
+Returns the list of all local per-replica values contained in `value`.
+
+Note: This only returns values on the worker initiated by this client. When
+using a `tf.distribute.Strategy` like
+`tf.distribute.experimental.MultiWorkerMirroredStrategy`, each worker will be
+its own client, and this function will only return values computed on that
+worker.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+value
+
+A value returned by experimental_run(), run(), or a variable
+created in scope`.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of values contained in value where ith element corresponds to
+ith replica. If value represents a single value, this returns
+(value,).
+

+
+
gather

+
+
+gather(
+    value, axis
+)
+

+
+Gather `value` across replicas along `axis` to the current device.
+
+Given a `tf.distribute.DistributedValues` or `tf.Tensor`-like object `value`,
+this API gathers and concatenates `value` across replicas along the `axis`-th
+dimension. The result is copied to the "current" device, which would typically
+be the CPU of the worker on which the program is running. For
+`tf.distribute.TPUStrategy`, it is the first TPU host. For multi-client
+`tf.distribute.MultiWorkerMirroredStrategy`, this is the CPU of each worker.
+
+This API can only be called in the cross-replica context. For a counterpart in
+the replica context, see `tf.distribute.ReplicaContext.all_gather`.
+
+Note: For all strategies except `tf.distribute.TPUStrategy`, the input `value`
+on different replicas must have the same rank, and their shapes must be the same
+in all dimensions except the `axis`-th dimension. In other words, their shapes
+cannot be different in a dimension `d` where `d` does not equal to the `axis`
+argument. For example, given a `tf.distribute.DistributedValues` with component
+tensors of shape `(1, 2, 3)` and `(1, 3, 3)` on two replicas, you can call
+`gather(..., axis=1, ...)` on it, but not `gather(..., axis=0, ...)` or
+`gather(..., axis=2, ...)`. However, for `tf.distribute.TPUStrategy.gather`, all
+tensors must have exactly the same rank and same shape.
+
+Note: Given a `tf.distribute.DistributedValues` `value`, its component tensors
+must have a non-zero rank. Otherwise, consider using `tf.expand_dims` before
+gathering them.
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+>>> # A DistributedValues with component tensor of shape (2, 1) on each replica
+... distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
+>>> @tf.function
+... def run():
+...   return strategy.gather(distributed_values, axis=0)
+>>> run()
+
+```
+
+Consider the following example for more combinations:
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3"])
+>>> single_tensor = tf.reshape(tf.range(6), shape=(1,2,3))
+>>> distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(single_tensor))
+>>> @tf.function
+... def run(axis):
+...   return strategy.gather(distributed_values, axis=axis)
+>>> axis=0
+>>> run(axis)
+
+>>> axis=1
+>>> run(axis)
+
+>>> axis=2
+>>> run(axis)
+
+```
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+value
+
+a tf.distribute.DistributedValues instance, e.g. returned by
+Strategy.run, to be combined into a single tensor. It can also be a
+regular tensor when used with tf.distribute.OneDeviceStrategy or the
+default strategy. The tensors that constitute the DistributedValues
+can only be dense tensors with non-zero rank, NOT a tf.IndexedSlices.
+

+axis
+
+0-D int32 Tensor. Dimension along which to gather. Must be in the
+range [0, rank(value)).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A Tensor that's the concatenation of value across replicas along
+axis dimension.
+

+
+
reduce

+
+
+reduce(
+    reduce_op, value, axis
+)
+

+
+Reduce `value` across replicas and return result on current device.
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+>>> def step_fn():
+...   i = tf.distribute.get_replica_context().replica_id_in_sync_group
+...   return tf.identity(i)
+>>>
+>>> per_replica_result = strategy.run(step_fn)
+>>> total = strategy.reduce("SUM", per_replica_result, axis=None)
+>>> total
+
+```
+
+To see how this would look with multiple replicas, consider the same example
+with MirroredStrategy with 2 GPUs:
+
+```python
+strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
+def step_fn():
+  i = tf.distribute.get_replica_context().replica_id_in_sync_group
+  return tf.identity(i)
+
+per_replica_result = strategy.run(step_fn)
+# Check devices on which per replica result is:
+strategy.experimental_local_results(per_replica_result)[0].device
+# /job:localhost/replica:0/task:0/device:GPU:0
+strategy.experimental_local_results(per_replica_result)[1].device
+# /job:localhost/replica:0/task:0/device:GPU:1
+
+total = strategy.reduce("SUM", per_replica_result, axis=None)
+# Check device on which reduced result is:
+total.device
+# /job:localhost/replica:0/task:0/device:CPU:0
+
+```
+
+This API is typically used for aggregating the results returned from different
+replicas, for reporting etc. For example, loss computed from different replicas
+can be averaged using this API before printing.
+
+Note: The result is copied to the "current" device - which would typically be
+the CPU of the worker on which the program is running. For `TPUStrategy`, it is
+the first TPU host. For multi client `MultiWorkerMirroredStrategy`, this is CPU
+of each worker.
+
+There are a number of different tf.distribute APIs for reducing values across
+replicas: * `tf.distribute.ReplicaContext.all_reduce`: This differs from
+`Strategy.reduce` in that it is for replica context and does not copy the
+results to the host device. `all_reduce` should be typically used for reductions
+inside the training step such as gradients. *
+`tf.distribute.StrategyExtended.reduce_to` and
+`tf.distribute.StrategyExtended.batch_reduce_to`: These APIs are more advanced
+versions of `Strategy.reduce` as they allow customizing the destination of the
+result. They are also called in cross replica context.
+
+*What should axis be?*
+
+Given a per-replica value returned by `run`, say a per-example loss, the batch
+will be divided across all the replicas. This function allows you to aggregate
+across replicas and optionally also across batch elements by specifying the axis
+parameter accordingly.
+
+For example, if you have a global batch size of 8 and 2 replicas, values for
+examples `[0, 1, 2, 3]` will be on replica 0 and `[4, 5, 6, 7]` will be on
+replica 1. With `axis=None`, `reduce` will aggregate only across replicas,
+returning `[0+4, 1+5, 2+6, 3+7]`. This is useful when each replica is computing
+a scalar or some other value that doesn't have a "batch" dimension (like a
+gradient or loss). `strategy.reduce("sum", per_replica_result, axis=None)`
+
+Sometimes, you will want to aggregate across both the global batch *and* all
+replicas. You can get this behavior by specifying the batch dimension as the
+`axis`, typically `axis=0`. In this case it would return a scalar
+`0+1+2+3+4+5+6+7`. `strategy.reduce("sum", per_replica_result, axis=0)`
+
+If there is a last partial batch, you will need to specify an axis so that the
+resulting shape is consistent across replicas. So if the last batch has size 6
+and it is divided into [0, 1, 2, 3] and [4, 5], you would get a shape mismatch
+unless you specify `axis=0`. If you specify `tf.distribute.ReduceOp.MEAN`, using
+`axis=0` will use the correct denominator of 6. Contrast this with computing
+`reduce_mean` to get a scalar value on each replica and this function to average
+those means, which will weigh some values `1/8` and others `1/4`.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+reduce_op
+
+a tf.distribute.ReduceOp value specifying how values should
+be combined. Allows using string representation of the enum such as
+"SUM", "MEAN".
+

+value
+
+a tf.distribute.DistributedValues instance, e.g. returned by
+Strategy.run, to be combined into a single tensor. It can also be a
+regular tensor when used with OneDeviceStrategy or default strategy.
+

+axis
+
+specifies the dimension to reduce along within each
+replica's tensor. Should typically be set to the batch dimension, or
+None to only reduce across replicas (e.g. if the tensor has no batch
+dimension).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A Tensor.
+

+
+
run

+
+
+run(
+    fn, args=(), kwargs=None, options=None
+)
+

+
+Invokes `fn` on each replica, with the given arguments.
+
+This method is the primary way to distribute your computation with a
+tf.distribute object. It invokes `fn` on each replica. If `args` or `kwargs`
+have `tf.distribute.DistributedValues`, such as those produced by a
+`tf.distribute.DistributedDataset` from
+`tf.distribute.Strategy.experimental_distribute_dataset` or
+`tf.distribute.Strategy.distribute_datasets_from_function`, when `fn` is
+executed on a particular replica, it will be executed with the component of
+`tf.distribute.DistributedValues` that correspond to that replica.
+
+`fn` is invoked under a replica context. `fn` may call
+`tf.distribute.get_replica_context()` to access members such as `all_reduce`.
+Please see the module-level docstring of tf.distribute for the concept of
+replica context.
+
+All arguments in `args` or `kwargs` can be a nested structure of tensors, e.g. a
+list of tensors, in which case `args` and `kwargs` will be passed to the `fn`
+invoked on each replica. Or `args` or `kwargs` can be
+`tf.distribute.DistributedValues` containing tensors or composite tensors, i.e.
+`tf.compat.v1.TensorInfo.CompositeTensor`, in which case each `fn` call will get
+the component of a `tf.distribute.DistributedValues` corresponding to its
+replica. Note that arbitrary Python values that are not of the types above are
+not supported.
+
+IMPORTANT: Depending on the implementation of `tf.distribute.Strategy` and
+whether eager execution is enabled, `fn` may be called one or more times. If
+`fn` is annotated with `tf.function` or `tf.distribute.Strategy.run` is called
+inside a `tf.function` (eager execution is disabled inside a `tf.function` by
+default), `fn` is called once per replica to generate a Tensorflow graph, which
+will then be reused for execution with new inputs. Otherwise, if eager execution
+is enabled, `fn` will be called once per replica every step just like regular
+python code.
+
+#### Example usage:
+
+1.  Constant tensor input.
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> tensor_input = tf.constant(3.0)
+    >>> @tf.function
+    ... def replica_fn(input):
+    ...   return input*2.0
+    >>> result = strategy.run(replica_fn, args=(tensor_input,))
+    >>> result
+    PerReplica:{
+      0: ,
+      1: 
+    }
+    ```
+
+2.  DistributedValues input. {: value=2}
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> @tf.function
+    ... def run():
+    ...   def value_fn(value_context):
+    ...     return value_context.num_replicas_in_sync
+    ...   distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...       value_fn))
+    ...   def replica_fn2(input):
+    ...     return input*2
+    ...   return strategy.run(replica_fn2, args=(distributed_values,))
+    >>> result = run()
+    >>> result
+    
+    ```
+
+3.  Use `tf.distribute.ReplicaContext` to allreduce values. {: value=3}
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["gpu:0", "gpu:1"])
+    >>> @tf.function
+    ... def run():
+    ...    def value_fn(value_context):
+    ...      return tf.constant(value_context.replica_id_in_sync_group)
+    ...    distributed_values = (
+    ...        strategy.experimental_distribute_values_from_function(
+    ...            value_fn))
+    ...    def replica_fn(input):
+    ...      return tf.distribute.get_replica_context().all_reduce(
+    ...          "sum", input)
+    ...    return strategy.run(replica_fn, args=(distributed_values,))
+    >>> result = run()
+    >>> result
+    PerReplica:{
+      0: ,
+      1: 
+    }
+    ```
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+fn
+
+The function to run on each replica.
+

+args
+
+Optional positional arguments to fn. Its element can be a tensor,
+a nested structure of tensors or a tf.distribute.DistributedValues.
+

+kwargs
+
+Optional keyword arguments to fn. Its element can be a tensor,
+a nested structure of tensors or a tf.distribute.DistributedValues.
+

+options
+
+An optional instance of tf.distribute.RunOptions specifying
+the options to run fn.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+Merged return value of fn across replicas. The structure of the return
+value is the same as the return value from fn. Each element in the
+structure can either be tf.distribute.DistributedValues, Tensor
+objects, or Tensors (for example, if running on a single replica).
+

+
+
scope

+
+
+scope()
+

+
+Context manager to make the strategy current and distribute variables.
+
+This method returns a context manager, and is used as follows:
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+>>> # Variable created inside scope:
+>>> with strategy.scope():
+...   mirrored_variable = tf.Variable(1.)
+>>> mirrored_variable
+MirroredVariable:{
+  0: ,
+  1: 
+}
+>>> # Variable created outside scope:
+>>> regular_variable = tf.Variable(1.)
+>>> regular_variable
+
+```
+
+*What happens when Strategy.scope is entered?*
+
+*   `strategy` is installed in the global context as the "current" strategy.
+    Inside this scope, `tf.distribute.get_strategy()` will now return this
+    strategy. Outside this scope, it returns the default no-op strategy.
+*   Entering the scope also enters the "cross-replica context". See
+    `tf.distribute.StrategyExtended` for an explanation on cross-replica and
+    replica contexts.
+*   Variable creation inside `scope` is intercepted by the strategy. Each
+    strategy defines how it wants to affect the variable creation. Sync
+    strategies like `MirroredStrategy`, `TPUStrategy` and
+    `MultiWorkerMiroredStrategy` create variables replicated on each replica,
+    whereas `ParameterServerStrategy` creates variables on the parameter
+    servers. This is done using a custom `tf.variable_creator_scope`.
+*   In some strategies, a default device scope may also be entered: in
+    `MultiWorkerMiroredStrategy`, a default device scope of "/CPU:0" is entered
+    on each worker.
+
+Note: Entering a scope does not automatically distribute a computation, except
+in the case of high level training framework like keras `model.fit`. If you're
+not using `model.fit`, you need to use `strategy.run` API to explicitly
+distribute that computation. See an example in the
+[custom training loop tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training).
+
+*What should be in scope and what should be outside?*
+
+There are a number of requirements on what needs to happen inside the scope.
+However, in places where we have information about which strategy is in use, we
+often enter the scope for the user, so they don't have to do it explicitly (i.e.
+calling those either inside or outside the scope is OK).
+
+*   Anything that creates variables that should be distributed variables must be
+    called in a `strategy.scope`. This can be accomplished either by directly
+    calling the variable creating function within the scope context, or by
+    relying on another API like `strategy.run` or `keras.Model.fit` to
+    automatically enter it for you. Any variable that is created outside scope
+    will not be distributed and may have performance implications. Some common
+    objects that create variables in TF are Models, Optimizers, Metrics. Such
+    objects should always be initialized in the scope, and any functions that
+    may lazily create variables (e.g., `Model.__call__()`, tracing a
+    `tf.function`, etc.) should similarly be called within scope. Another source
+    of variable creation can be a checkpoint restore - when variables are
+    created lazily. Note that any variable created inside a strategy captures
+    the strategy information. So reading and writing to these variables outside
+    the `strategy.scope` can also work seamlessly, without the user having to
+    enter the scope.
+*   Some strategy APIs (such as `strategy.run` and `strategy.reduce`) which
+    require to be in a strategy's scope, enter the scope automatically, which
+    means when using those APIs you don't need to explicitly enter the scope
+    yourself.
+*   When a `tf.keras.Model` is created inside a `strategy.scope`, the Model
+    object captures the scope information. When high level training framework
+    methods such as `model.compile`, `model.fit`, etc. are then called, the
+    captured scope will be automatically entered, and the associated strategy
+    will be used to distribute the training etc. See a detailed example in
+    [distributed keras tutorial](https://www.tensorflow.org/tutorials/distribute/keras).
+    WARNING: Simply calling `model(..)` does not automatically enter the
+    captured scope -- only high level training framework APIs support this
+    behavior: `model.compile`, `model.fit`, `model.evaluate`, `model.predict`
+    and `model.save` can all be called inside or outside the scope.
+*   The following can be either inside or outside the scope:
+    *   Creating the input datasets
+    *   Defining `tf.function`s that represent your training step
+    *   Saving APIs such as `tf.saved_model.save`. Loading creates variables, so
+        that should go inside the scope if you want to train the model in a
+        distributed way.
+    *   Checkpoint saving. As mentioned above - `checkpoint.restore` may
+        sometimes need to be inside scope if it creates variables.
+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A context manager.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/PassthruDatasetProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/PassthruDatasetProvider.md
new file mode 100644
index 00000000..3e1bdaea
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/PassthruDatasetProvider.md
@@ -0,0 +1,40 @@
+# runner.PassthruDatasetProvider
+
+
+
+
+ View source
+on GitHub 
+
+Builds a `tf.data.Dataset` from a pass thru dataset.
+
+Inherits From: [`DatasetProvider`](../runner/DatasetProvider.md)
+
+
+runner.PassthruDatasetProvider(
+    dataset: tf.data.Dataset,
+    *,
+    shuffle_datasets: bool = False,
+    examples_shuffle_size: Optional[int] = None
+)
+

+
+
+
+Passes any `dataset` thru: omitting any sharding. For detailed documentation,
+see the filename dataset provider complement: `SimpleDatasetsProvider.`
+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+get_dataset(
+    _: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Gets a `tf.data.Dataset` omitting any input context.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/PassthruSampleDatasetsProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/PassthruSampleDatasetsProvider.md
new file mode 100644
index 00000000..0c841412
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/PassthruSampleDatasetsProvider.md
@@ -0,0 +1,46 @@
+# runner.PassthruSampleDatasetsProvider
+
+
+
+
+ View source
+on GitHub 
+
+Builds a sampled `tf.data.Dataset` from multiple pass thru datasets.
+
+Inherits From: [`DatasetProvider`](../runner/DatasetProvider.md)
+
+
+runner.PassthruSampleDatasetsProvider(
+    principal_dataset: tf.data.Dataset,
+    extra_datasets: Sequence[tf.data.Dataset],
+    principal_weight: Optional[float] = None,
+    extra_weights: Optional[Sequence[float]] = None,
+    *,
+    principal_cardinality: Optional[int] = None,
+    fixed_cardinality: bool = False,
+    shuffle_dataset: bool = False,
+    examples_shuffle_size: Optional[int] = None
+)
+

+
+
+
+Passes any `principal_dataset` and `extra_datasets` thru: omitting any sharding.
+For detailed documentation, see the filename dataset provider complement:
+`SimpleSampleDatasetsProvider.`
+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+get_dataset(
+    _: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Gets a sampled `tf.data.Dataset` omitting any input context.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/Predictions.md b/tensorflow_gnn/docs/api_docs/python/runner/Predictions.md
new file mode 100644
index 00000000..31446b86
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/Predictions.md
@@ -0,0 +1,17 @@
+# runner.Predictions
+
+
+
+This symbol is a **type alias**.
+
+#### Source:
+
+
+Predictions = Union[
+    tf.Tensor,
+    tf.RaggedTensor,
+    Mapping[str, Union[tf.Tensor, tf.RaggedTensor]]
+]
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeBinaryClassification.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeBinaryClassification.md
new file mode 100644
index 00000000..817ae0f5
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeBinaryClassification.md
@@ -0,0 +1,216 @@
+# runner.RootNodeBinaryClassification
+
+
+
+
+ View source
+on GitHub 
+
+Root node binary (or multi-label) classification.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeBinaryClassification(
+    node_set_name: str,
+    units: int = 1,
+    *,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    name: str = 'classification_logits',
+    label_fn: Optional[LabelFn] = None,
+    label_feature_name: Optional[str] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+node_set_name
+
+The node set containing the root node.
+

+units
+
+The units for the classification head. (Typically 1 for binary
+classification and the number of labels for multi-label classification.)
+

+state_name
+
+The feature name for activations (e.g.: tfgnn.HIDDEN_STATE).
+

+name
+
+The classification head's layer name. To control the naming of saved
+model outputs see the runner model exporters (e.g.,
+KerasModelExporter).
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> Field
+

+
+Gather activations from root nodes.
+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+Returns arbitrary task specific losses.
+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Returns arbitrary task specific metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for classification.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for classification.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The classification logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeLabelFn.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeLabelFn.md
new file mode 100644
index 00000000..80a27eee
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeLabelFn.md
@@ -0,0 +1,20 @@
+# runner.RootNodeLabelFn
+
+
+
+
+ View source
+on GitHub 
+
+Reads out a `tfgnn.Field` from the `GraphTensor` root (i.e. first) node.
+
+
+runner.RootNodeLabelFn(
+    node_set_name: tfgnn.NodeSetName,
+    *,
+    feature_name: tfgnn.FieldName = tfgnn.HIDDEN_STATE,
+    **kwargs
+)
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteError.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteError.md
new file mode 100644
index 00000000..08daf609
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteError.md
@@ -0,0 +1,194 @@
+# runner.RootNodeMeanAbsoluteError
+
+
+
+
+ View source
+on GitHub 
+
+Mean absolute error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMeanAbsoluteError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteLogarithmicError.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteLogarithmicError.md
new file mode 100644
index 00000000..57c56240
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsoluteLogarithmicError.md
@@ -0,0 +1,152 @@
+# runner.RootNodeMeanAbsoluteLogarithmicError
+
+
+
+
+ View source
+on GitHub 
+
+Root node mean absolute logarithmic error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMeanAbsoluteLogarithmicError(
+    reduction: tf.keras.losses.Reduction = AUTO,
+    name: Optional[str] = None,
+    **kwargs
+)
+

+
+
+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a head with ReLU for nonnegative regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor use for prediction.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The nonnegative logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsolutePercentageError.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsolutePercentageError.md
new file mode 100644
index 00000000..defda629
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanAbsolutePercentageError.md
@@ -0,0 +1,194 @@
+# runner.RootNodeMeanAbsolutePercentageError
+
+
+
+
+ View source
+on GitHub 
+
+Mean absolute percentage error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMeanAbsolutePercentageError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredError.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredError.md
new file mode 100644
index 00000000..18b717d5
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredError.md
@@ -0,0 +1,194 @@
+# runner.RootNodeMeanSquaredError
+
+
+
+
+ View source
+on GitHub 
+
+Mean squared error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMeanSquaredError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogScaledError.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogScaledError.md
new file mode 100644
index 00000000..d90fc77c
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogScaledError.md
@@ -0,0 +1,155 @@
+# runner.RootNodeMeanSquaredLogScaledError
+
+
+
+
+ View source
+on GitHub 
+
+Mean squared log scaled error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMeanSquaredLogScaledError(
+    *args,
+    alpha_loss_param: float = 5.0,
+    epsilon_loss_param: float = 1e-08,
+    reduction: tf.keras.losses.Reduction = AUTO,
+    name: Optional[str] = None,
+    **kwargs
+)
+

+
+
+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogarithmicError.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogarithmicError.md
new file mode 100644
index 00000000..bbe8e2d8
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMeanSquaredLogarithmicError.md
@@ -0,0 +1,194 @@
+# runner.RootNodeMeanSquaredLogarithmicError
+
+
+
+
+ View source
+on GitHub 
+
+Mean squared logarithmic error task.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMeanSquaredLogarithmicError(
+    node_set_name: str,
+    *,
+    units: int = 1,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    **kwargs
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+units
+
+The units for the regression head.
+

+name
+
+The regression head's layer name. This name typically appears in
+the exported model's SignatureDef.
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> tf.Tensor
+

+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Regression metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for regression.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for regression.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The regression logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMulticlassClassification.md b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMulticlassClassification.md
new file mode 100644
index 00000000..a87e6a36
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RootNodeMulticlassClassification.md
@@ -0,0 +1,233 @@
+# runner.RootNodeMulticlassClassification
+
+
+
+
+ View source
+on GitHub 
+
+Root node multiclass classification.
+
+Inherits From: [`Task`](../runner/Task.md)
+
+
+runner.RootNodeMulticlassClassification(
+    node_set_name: str,
+    *,
+    num_classes: Optional[int] = None,
+    class_names: Optional[Sequence[str]] = None,
+    per_class_statistics: bool = False,
+    state_name: str = tfgnn.HIDDEN_STATE,
+    name: str = 'classification_logits',
+    label_fn: Optional[LabelFn] = None,
+    label_feature_name: Optional[str] = None
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+node_set_name
+
+The node set containing the root node.
+

+num_classes
+
+The number of classes. Exactly one of num_classes or
+class_names must be specified
+

+class_names
+
+The class names. Exactly one of num_classes or
+class_names must be specified
+

+per_class_statistics
+
+Whether to compute statistics per class.
+

+state_name
+
+The feature name for activations (e.g.: tfgnn.HIDDEN_STATE).
+

+name
+
+The classification head's layer name. To control the naming of saved
+model outputs see the runner model exporters (e.g.,
+KerasModelExporter).
+

+label_fn
+
+A label extraction function. This function mutates the input
+GraphTensor. Mutually exclusive with label_feature_name.
+

+label_feature_name
+
+A label feature name for readout from the auxiliary
+'_readout' node set. Readout does not mutate the input GraphTensor.
+Mutually exclusive with label_fn.
+

+
+## Methods
+
+
gather_activations

+
+View
+source
+
+
+gather_activations(
+    inputs: GraphTensor
+) -> Field
+

+
+Gather activations from root nodes.
+
+
losses

+
+View
+source
+
+
+losses() -> interfaces.Losses
+

+
+Sparse categorical crossentropy loss.
+
+
metrics

+
+View
+source
+
+
+metrics() -> interfaces.Metrics
+

+
+Sparse categorical metrics.
+
+
predict

+
+View
+source
+
+
+predict(
+    inputs: tfgnn.GraphTensor
+) -> interfaces.Predictions
+

+
+Apply a linear head for classification.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A tfgnn.GraphTensor for classification.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The classification logits.
+

+
+
preprocess

+
+View
+source
+
+
+preprocess(
+    inputs: GraphTensor
+) -> tuple[GraphTensor, Field]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/RunResult.md b/tensorflow_gnn/docs/api_docs/python/runner/RunResult.md
new file mode 100644
index 00000000..8e8aa163
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/RunResult.md
@@ -0,0 +1,71 @@
+# runner.RunResult
+
+
+
+
+ View source
+on GitHub 
+
+Holds the return values of `run(...)`.
+
+
+runner.RunResult(
+    preprocess_model: Optional[tf.keras.Model],
+    base_model: tf.keras.Model,
+    trained_model: tf.keras.Model
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+preprocess_model
+
+Keras model containing only the computation for
+preprocessing inputs. It is not trained. The model takes serialized
+GraphTensors as its inputs and returns preprocessed GraphTensors.
+None when no preprocess model exists.
+

+base_model
+
+Keras base GNN (as returned by the user provided model_fn).
+The model both takes and returns GraphTensors. The model contains
+any--but not all--trained weights. The trained_model contains all
+base_model trained weights in addition to any prediction trained
+weights.
+

+trained_model
+
+Keras model for the e2e GNN. (Base GNN plus any prediction
+head(s).) The model takes preprocess_model output as its inputs and
+returns Task predictions as its output. Output matches the structure of
+the Task: an atom for single- or a mapping for multi- Task training.
+The model contains all trained weights.
+

+
+## Methods
+
+
__eq__

+
+
+__eq__(
+    other
+)
+

+
+Return self==value.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/SampleTFRecordDatasetsProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/SampleTFRecordDatasetsProvider.md
new file mode 100644
index 00000000..d3950aec
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/SampleTFRecordDatasetsProvider.md
@@ -0,0 +1,229 @@
+# runner.SampleTFRecordDatasetsProvider
+
+
+
+
+ View source
+on GitHub 
+
+Builds a sampling `tf.data.Dataset` from multiple filenames.
+
+Inherits From:
+[`SimpleSampleDatasetsProvider`](../runner/SimpleSampleDatasetsProvider.md),
+[`DatasetProvider`](../runner/DatasetProvider.md)
+
+
+runner.SampleTFRecordDatasetsProvider(
+    *args, **kwargs
+)
+

+
+
+
+For complete explanations regarding sampling see `_process_sampled_dataset()`.
+
+This `SimpleSampleDatasetsProvider` builds a `tf.data.Dataset` as follows:
+
+-   The object is initialized with a list of filenames specified by
+    `principle_filenames` and `extra_filenames` argument. For convenience, the
+    corresponding file pattern `principal_file_pattern` and
+    `extra_file_patterns` can be specified instead, which will be expanded to a
+    sorted list.
+-   The filenames are sharded between replicas according to the `InputContext`
+    (order matters).
+-   Filenames are shuffled per replica (if requested).
+-   Examples from all file patterns are sampled according to `principal_weight`
+    and `extra_weights.`
+-   The files in each shard are interleaved after being read by the
+    `interleave_fn`.
+-   Examples are shuffled (if requested), auto-prefetched, and returned for use
+    in one replica of the trainer.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+principal_file_pattern
+
+A principal file pattern for sampling, to be
+expanded by tf.io.gfile.glob and sorted into the list of
+principal_filenames.
+

+extra_file_patterns
+
+File patterns, to be expanded by tf.io.gfile.glob
+and sorted into the list of extra_filenames.
+

+principal_weight
+
+An optional weight for the dataset corresponding to
+principal_file_pattern. Required iff extra_weights are also
+provided.
+

+extra_weights
+
+Optional weights corresponding to file_patterns for
+sampling. Required iff principal_weight is also provided.
+

+principal_filenames
+
+A list of principal filenames, specified explicitly.
+This argument is mutually exclusive with principal_file_pattern.
+

+extra_filenames
+
+A list of extra filenames, specified explicitly.
+This argument is mutually exclusive with extra_file_patterns.
+

+principal_cardinality
+
+Iff fixed_cardinality=True, the size of the
+returned dataset is computed as principal_cardinality /
+principal_weight (with a default of uniform weights).
+

+fixed_cardinality
+
+Whether to take a fixed number of elements.
+

+shuffle_filenames
+
+If enabled, filenames will be shuffled after sharding
+ between replicas, before any file reads. Through interleaving, some
+files may be read in parallel: the details are auto-tuned for throughput.
+

+interleave_fn
+
+A fn applied with tf.data.Dataset.interleave.
+

+examples_shuffle_size
+
+An optional buffer size for example shuffling. If
+specified, the size is adjusted to shuffle_size //
+(len(file_patterns) + 1).
+

+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+get_dataset(
+    context: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Creates a `tf.data.Dataset` by sampling.
+
+The contents of the resulting `tf.data.Dataset` are sampled from several
+sources, each stored as a sharded dataset: * one principal input, whose size
+determines the size of the resulting `tf.data.Dataset`; * zero or more side
+inputs, which are repeated if necessary to preserve the requested samping
+weights.
+
+Each input dataset is shared before interleaving. The result of interleaving is
+only shuffled if a `examples_shuffle_size` is provided.
+
+Datasets are sampled from with `tf.data.Dataset.sample_from_datasets.` For
+sampling details, please refer to the TensorFlow documentation at:
+https://www.tensorflow.org/api_docs/python/tf/data/Dataset#sample_from_datasets.
+
+Two methods are supported to determine the end of the resulting
+`tf.data.Dataset`:
+
+fixed_cardinality=True) Returns a dataset with a fixed cardinality, set at
+`principal_cardinality` // `principal_weight.` `principal_dataset` and
+`principal_cardinality` are required for this method. `principal_weight` is
+required iff `extra_weights` are also provided.
+
+fixed_cardinality=False) Returns a dataset that ends after the principal input
+has been exhausted, subject to the random selection of samples.
+`principal_dataset` is required for this method. `principal_weight` is required
+iff `extra_weights` are also provided.
+
+The choice of `principal_dataset` is important and should, in most cases, be
+chosen as the largest underlying dataset as compared to `extra_datasets.`
+`positives` and `negatives` where `len(negatives)` >> `len(positives)` and with
+`positives` corresponding to `principal_dataset,` the desired behavior of epochs
+determined by the exhaustion of `positives` and the continued mixing of unique
+elements from `negatives` may not occur: On sampled dataset reiteration
+`positives` will again be exhausted but elements from `negatives` may be those
+same seen in the previous epoch (as they occur at the beginning of the same,
+reiterated underlying `negatives` dataset). In this case, the recommendations
+are to:
+
+1) Reformulate the sampling in terms of the larger dataset (`negatives`), where,
+with `fixed_cardinality=False`, if the exhaustion of `negatives` is desired, or,
+with `fixed_cardinality=True`, when `principal_cardinality` can be used to
+specify the desired number of elements from `negatives.` 2) Ensure that the
+underlying `principal_dataset` of `negatives` are well-sharded. In this way, the
+nondeterminism of interleaving will randomly access elements of `negatives` on
+reiteration.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+context
+
+An tf.distribute.InputContext for sharding.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.data.Dataset.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/SimpleDatasetProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/SimpleDatasetProvider.md
new file mode 100644
index 00000000..bb730e58
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/SimpleDatasetProvider.md
@@ -0,0 +1,97 @@
+# runner.SimpleDatasetProvider
+
+
+
+
+ View source
+on GitHub 
+
+Builds a `tf.data.Dataset` from a list of files.
+
+Inherits From: [`DatasetProvider`](../runner/DatasetProvider.md)
+
+
+runner.SimpleDatasetProvider(
+    file_pattern: Optional[str] = None,
+    *,
+    filenames: Optional[Sequence[str]] = None,
+    shuffle_filenames: bool = False,
+    interleave_fn: Callable[..., tf.data.Dataset],
+    examples_shuffle_size: Optional[int] = None
+)
+

+
+
+
+This `SimpleDatasetProvider` builds a `tf.data.Dataset` as follows: - The object
+is initialized with a list of filenames. For convenience, a file pattern can be
+specified instead, which will be expanded to a sorted list. - The filenames are
+sharded between replicas according to the `InputContext` (order matters). -
+Filenames are shuffled per replica (if requested). - The files in each shard are
+interleaved after being read by the `interleave_fn`. - Examples are shuffled (if
+requested), auto-prefetched, and returned for use in one replica of the trainer.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+file_pattern
+
+A file pattern, to be expanded by tf.io.gfile.glob
+and sorted into the list of all filenames.
+

+filenames
+
+A list of all filenames, specified explicitly.
+This argument is mutually exclusive with file_pattern.
+

+shuffle_filenames
+
+If enabled, filenames will be shuffled after sharding
+between replicas, before any file reads. Through interleaving, some
+files may be read in parallel: the details are auto-tuned for
+throughput.
+

+interleave_fn
+
+A callback that receives a single filename and returns
+a tf.data.Dataset with the tf.Example values from that file.
+

+examples_shuffle_size
+
+An optional buffer size for example shuffling.
+

+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+get_dataset(
+    context: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Gets a `tf.data.Dataset` by `context` per replica.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/SimpleSampleDatasetsProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/SimpleSampleDatasetsProvider.md
new file mode 100644
index 00000000..87e5c5e4
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/SimpleSampleDatasetsProvider.md
@@ -0,0 +1,238 @@
+# runner.SimpleSampleDatasetsProvider
+
+
+
+
+ View source
+on GitHub 
+
+Builds a sampling `tf.data.Dataset` from multiple filenames.
+
+Inherits From: [`DatasetProvider`](../runner/DatasetProvider.md)
+
+
+runner.SimpleSampleDatasetsProvider(
+    principal_file_pattern: Optional[str] = None,
+    extra_file_patterns: Optional[Sequence[str]] = None,
+    principal_weight: Optional[float] = None,
+    extra_weights: Optional[Sequence[float]] = None,
+    *,
+    principal_filenames: Optional[Sequence[str]] = None,
+    extra_filenames: Optional[Sequence[Sequence[str]]] = None,
+    principal_cardinality: Optional[int] = None,
+    fixed_cardinality: bool = False,
+    shuffle_filenames: bool = False,
+    interleave_fn: Callable[..., tf.data.Dataset],
+    examples_shuffle_size: Optional[int] = None
+)
+

+
+
+
+For complete explanations regarding sampling see `_process_sampled_dataset()`.
+
+This `SimpleSampleDatasetsProvider` builds a `tf.data.Dataset` as follows:
+
+-   The object is initialized with a list of filenames specified by
+    `principle_filenames` and `extra_filenames` argument. For convenience, the
+    corresponding file pattern `principal_file_pattern` and
+    `extra_file_patterns` can be specified instead, which will be expanded to a
+    sorted list.
+-   The filenames are sharded between replicas according to the `InputContext`
+    (order matters).
+-   Filenames are shuffled per replica (if requested).
+-   Examples from all file patterns are sampled according to `principal_weight`
+    and `extra_weights.`
+-   The files in each shard are interleaved after being read by the
+    `interleave_fn`.
+-   Examples are shuffled (if requested), auto-prefetched, and returned for use
+    in one replica of the trainer.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+principal_file_pattern
+
+A principal file pattern for sampling, to be
+expanded by tf.io.gfile.glob and sorted into the list of
+principal_filenames.
+

+extra_file_patterns
+
+File patterns, to be expanded by tf.io.gfile.glob
+and sorted into the list of extra_filenames.
+

+principal_weight
+
+An optional weight for the dataset corresponding to
+principal_file_pattern. Required iff extra_weights are also
+provided.
+

+extra_weights
+
+Optional weights corresponding to file_patterns for
+sampling. Required iff principal_weight is also provided.
+

+principal_filenames
+
+A list of principal filenames, specified explicitly.
+This argument is mutually exclusive with principal_file_pattern.
+

+extra_filenames
+
+A list of extra filenames, specified explicitly.
+This argument is mutually exclusive with extra_file_patterns.
+

+principal_cardinality
+
+Iff fixed_cardinality=True, the size of the
+returned dataset is computed as principal_cardinality /
+principal_weight (with a default of uniform weights).
+

+fixed_cardinality
+
+Whether to take a fixed number of elements.
+

+shuffle_filenames
+
+If enabled, filenames will be shuffled after sharding
+ between replicas, before any file reads. Through interleaving, some
+files may be read in parallel: the details are auto-tuned for throughput.
+

+interleave_fn
+
+A fn applied with tf.data.Dataset.interleave.
+

+examples_shuffle_size
+
+An optional buffer size for example shuffling. If
+specified, the size is adjusted to shuffle_size //
+(len(file_patterns) + 1).
+

+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+get_dataset(
+    context: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Creates a `tf.data.Dataset` by sampling.
+
+The contents of the resulting `tf.data.Dataset` are sampled from several
+sources, each stored as a sharded dataset: * one principal input, whose size
+determines the size of the resulting `tf.data.Dataset`; * zero or more side
+inputs, which are repeated if necessary to preserve the requested samping
+weights.
+
+Each input dataset is shared before interleaving. The result of interleaving is
+only shuffled if a `examples_shuffle_size` is provided.
+
+Datasets are sampled from with `tf.data.Dataset.sample_from_datasets.` For
+sampling details, please refer to the TensorFlow documentation at:
+https://www.tensorflow.org/api_docs/python/tf/data/Dataset#sample_from_datasets.
+
+Two methods are supported to determine the end of the resulting
+`tf.data.Dataset`:
+
+fixed_cardinality=True) Returns a dataset with a fixed cardinality, set at
+`principal_cardinality` // `principal_weight.` `principal_dataset` and
+`principal_cardinality` are required for this method. `principal_weight` is
+required iff `extra_weights` are also provided.
+
+fixed_cardinality=False) Returns a dataset that ends after the principal input
+has been exhausted, subject to the random selection of samples.
+`principal_dataset` is required for this method. `principal_weight` is required
+iff `extra_weights` are also provided.
+
+The choice of `principal_dataset` is important and should, in most cases, be
+chosen as the largest underlying dataset as compared to `extra_datasets.`
+`positives` and `negatives` where `len(negatives)` >> `len(positives)` and with
+`positives` corresponding to `principal_dataset,` the desired behavior of epochs
+determined by the exhaustion of `positives` and the continued mixing of unique
+elements from `negatives` may not occur: On sampled dataset reiteration
+`positives` will again be exhausted but elements from `negatives` may be those
+same seen in the previous epoch (as they occur at the beginning of the same,
+reiterated underlying `negatives` dataset). In this case, the recommendations
+are to:
+
+1) Reformulate the sampling in terms of the larger dataset (`negatives`), where,
+with `fixed_cardinality=False`, if the exhaustion of `negatives` is desired, or,
+with `fixed_cardinality=True`, when `principal_cardinality` can be used to
+specify the desired number of elements from `negatives.` 2) Ensure that the
+underlying `principal_dataset` of `negatives` are well-sharded. In this way, the
+nondeterminism of interleaving will randomly access elements of `negatives` on
+reiteration.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+context
+
+An tf.distribute.InputContext for sharding.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.data.Dataset.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/SubmoduleExporter.md b/tensorflow_gnn/docs/api_docs/python/runner/SubmoduleExporter.md
new file mode 100644
index 00000000..bb2dc054
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/SubmoduleExporter.md
@@ -0,0 +1,125 @@
+# runner.SubmoduleExporter
+
+
+
+
+ View source
+on GitHub 
+
+Exports a Keras submodule.
+
+Inherits From: [`ModelExporter`](../runner/ModelExporter.md)
+
+
+runner.SubmoduleExporter(
+    sublayer_name: str,
+    *,
+    output_names: Optional[Any] = None,
+    subdirectory: Optional[str] = None,
+    include_preprocessing: bool = False,
+    options: Optional[tf.saved_model.SaveOptions] = None
+)
+

+
+
+
+Given a `RunResult`, this exporter creates and exports a submodule with inputs
+identical to the trained model and outputs from some intermediate layer (named
+`sublayer_name`). For example, with pseudocode:
+
+`trained_model = tf.keras.Sequential([layer1, layer2, layer3, layer4])` and
+`SubmoduleExporter(sublayer_name='layer2')`
+
+The exported submodule is:
+
+`submodule = tf.keras.Sequential([layer1, layer2])`
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+sublayer_name
+
+The name of the submodule's final layer.
+

+output_names
+
+The names for output Tensor(s), see: KerasModelExporter.
+

+subdirectory
+
+An optional subdirectory, if set: submodules are exported
+to os.path.join(export_dir, subdirectory).
+

+include_preprocessing
+
+Whether to include any preprocess_model.
+

+options
+
+Options for saving to a TensorFlow SavedModel.
+

+
+## Methods
+
+
save

+
+View
+source
+
+
+save(
+    run_result: runner.RunResult,
+    export_dir: str
+)
+

+
+Saves a Keras model submodule.
+
+Importantly: the `run_result.preprocess_model`, if provided, and
+`run_result.trained_model` are stacked before any export. Stacking involves the
+chaining of the first output of `run_result.preprocess_model` to the only input
+of `run_result.trained_model.` The result is a model with the input of
+`run_result.preprocess_model` and the output of `run_result.trained_model.`
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+run_result
+
+A RunResult from training.
+

+export_dir
+
+A destination directory.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/TFDataServiceConfig.md b/tensorflow_gnn/docs/api_docs/python/runner/TFDataServiceConfig.md
new file mode 100644
index 00000000..f94a5164
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/TFDataServiceConfig.md
@@ -0,0 +1,66 @@
+# runner.TFDataServiceConfig
+
+
+
+
+ View source
+on GitHub 
+
+Provides tf.data service related configuration options.
+
+
+runner.TFDataServiceConfig(
+    tf_data_service_address: str,
+    tf_data_service_job_name: str,
+    tf_data_service_mode: Union[str, tf.data.experimental.service.ShardingPolicy]
+)
+

+
+
+
+tf.data service has data flexible visitation guarantees, its impact over your
+training pipelines will be empirical. Check out the tf.data service internals
+and operation details from
+https://www.tensorflow.org/api_docs/python/tf/data/experimental/service.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+tf_data_service_address
+
+Dataclass field
+

+tf_data_service_job_name
+
+Dataclass field
+

+tf_data_service_mode
+
+Dataclass field
+

+
+## Methods
+
+
__eq__

+
+
+__eq__(
+    other
+)
+

+
+Return self==value.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/TFRecordDatasetProvider.md b/tensorflow_gnn/docs/api_docs/python/runner/TFRecordDatasetProvider.md
new file mode 100644
index 00000000..d3b381de
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/TFRecordDatasetProvider.md
@@ -0,0 +1,93 @@
+# runner.TFRecordDatasetProvider
+
+
+
+
+ View source
+on GitHub 
+
+Builds a `tf.data.Dataset` from a list of files.
+
+Inherits From: [`SimpleDatasetProvider`](../runner/SimpleDatasetProvider.md),
+[`DatasetProvider`](../runner/DatasetProvider.md)
+
+
+runner.TFRecordDatasetProvider(
+    *args, **kwargs
+)
+

+
+
+
+This `SimpleDatasetProvider` builds a `tf.data.Dataset` as follows: - The object
+is initialized with a list of filenames. For convenience, a file pattern can be
+specified instead, which will be expanded to a sorted list. - The filenames are
+sharded between replicas according to the `InputContext` (order matters). -
+Filenames are shuffled per replica (if requested). - The files in each shard are
+interleaved after being read by the `interleave_fn`. - Examples are shuffled (if
+requested), auto-prefetched, and returned for use in one replica of the trainer.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+file_pattern
+
+A file pattern, to be expanded by tf.io.gfile.glob
+and sorted into the list of all filenames.
+

+filenames
+
+A list of all filenames, specified explicitly.
+This argument is mutually exclusive with file_pattern.
+

+shuffle_filenames
+
+If enabled, filenames will be shuffled after sharding
+between replicas, before any file reads. Through interleaving, some
+files may be read in parallel: the details are auto-tuned for
+throughput.
+

+interleave_fn
+
+A callback that receives a single filename and returns
+a tf.data.Dataset with the tf.Example values from that file.
+

+examples_shuffle_size
+
+An optional buffer size for example shuffling.
+

+
+## Methods
+
+
get_dataset

+
+View
+source
+
+
+get_dataset(
+    context: tf.distribute.InputContext
+) -> tf.data.Dataset
+

+
+Gets a `tf.data.Dataset` by `context` per replica.
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/TPUStrategy.md b/tensorflow_gnn/docs/api_docs/python/runner/TPUStrategy.md
new file mode 100644
index 00000000..d80ce7b5
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/TPUStrategy.md
@@ -0,0 +1,1228 @@
+# runner.TPUStrategy
+
+
+
+
+ View source
+on GitHub 
+
+A `TPUStrategy` convenience wrapper.
+
+
+runner.TPUStrategy(
+    tpu: str = ''
+)
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+tpu_cluster_resolver
+
+A
+tf.distribute.cluster_resolver.TPUClusterResolver instance, which
+provides information about the TPU cluster. If None, it will assume
+running on a local TPU worker.
+

+experimental_device_assignment
+
+Optional
+tf.tpu.experimental.DeviceAssignment to specify the placement of
+replicas on the TPU cluster.
+

+experimental_spmd_xla_partitioning
+
+If True, enable the SPMD (Single
+Program Multiple Data) mode in XLA compiler. This flag only affects the
+performance of XLA compilation and the HBM requirement of the compiled
+TPU program. Ceveat: if this flag is True, calling
+tf.distribute.TPUStrategy.experimental_assign_to_logical_device will
+result in a ValueError.
+

+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
Attributes
  cluster_resolver  
+Returns the cluster resolver associated with this strategy.
+
+tf.distribute.TPUStrategy provides the associated
+tf.distribute.cluster_resolver.ClusterResolver. If the user provides one
+in **init**, that instance is returned; if the user does not, a default
+tf.distribute.cluster_resolver.TPUClusterResolver is provided.
+

+extended
+
+tf.distribute.StrategyExtended with additional methods.
+

+num_replicas_in_sync
+
+Returns number of replicas over which gradients are aggregated.
+

+
+## Methods
+
+
distribute_datasets_from_function

+
+
+distribute_datasets_from_function(
+    dataset_fn, options=None
+)
+

+
+Distributes `tf.data.Dataset` instances created by calls to `dataset_fn`.
+
+The argument `dataset_fn` that users pass in is an input function that has a
+`tf.distribute.InputContext` argument and returns a `tf.data.Dataset` instance.
+It is expected that the returned dataset from `dataset_fn` is already batched by
+per-replica batch size (i.e. global batch size divided by the number of replicas
+in sync) and sharded. `tf.distribute.Strategy.distribute_datasets_from_function`
+does not batch or shard the `tf.data.Dataset` instance returned from the input
+function. `dataset_fn` will be called on the CPU device of each of the workers
+and each generates a dataset where every replica on that worker will dequeue one
+batch of inputs (i.e. if a worker has two replicas, two batches will be dequeued
+from the `Dataset` every step).
+
+This method can be used for several purposes. First, it allows you to specify
+your own batching and sharding logic. (In contrast,
+`tf.distribute.experimental_distribute_dataset` does batching and sharding for
+you.) For example, where `experimental_distribute_dataset` is unable to shard
+the input files, this method might be used to manually shard the dataset
+(avoiding the slow fallback behavior in `experimental_distribute_dataset`). In
+cases where the dataset is infinite, this sharding can be done by creating
+dataset replicas that differ only in their random seed.
+
+The `dataset_fn` should take an `tf.distribute.InputContext` instance where
+information about batching and input replication can be accessed.
+
+You can use `element_spec` property of the `tf.distribute.DistributedDataset`
+returned by this API to query the `tf.TypeSpec` of the elements returned by the
+iterator. This can be used to set the `input_signature` property of a
+`tf.function`. Follow `tf.distribute.DistributedDataset.element_spec` to see an
+example.
+
+IMPORTANT: The `tf.data.Dataset` returned by `dataset_fn` should have a
+per-replica batch size, unlike `experimental_distribute_dataset`, which uses the
+global batch size. This may be computed using
+`input_context.get_per_replica_batch_size`.
+
+Note: If you are using TPUStrategy, the order in which the data is processed by
+the workers when using `tf.distribute.Strategy.experimental_distribute_dataset`
+or `tf.distribute.Strategy.distribute_datasets_from_function` is not guaranteed.
+This is typically required if you are using `tf.distribute` to scale prediction.
+You can however insert an index for each element in the batch and order outputs
+accordingly. Refer to
+[this snippet](https://www.tensorflow.org/tutorials/distribute/input#caveats)
+for an example of how to order outputs.
+
+Note: Stateful dataset transformations are currently not supported with
+`tf.distribute.experimental_distribute_dataset` or
+`tf.distribute.distribute_datasets_from_function`. Any stateful ops that the
+dataset may have are currently ignored. For example, if your dataset has a
+`map_fn` that uses `tf.random.uniform` to rotate an image, then you have a
+dataset graph that depends on state (i.e the random seed) on the local machine
+where the python process is being executed.
+
+For a tutorial on more usage and properties of this method, refer to the
+[tutorial on distributed input](https://www.tensorflow.org/tutorials/distribute/input#tfdistributestrategyexperimental_distribute_datasets_from_function)).
+If you are interested in last partial batch handling, read
+[this section](https://www.tensorflow.org/tutorials/distribute/input#partial_batches).
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+dataset_fn
+
+A function taking a tf.distribute.InputContext instance and
+returning a tf.data.Dataset.
+

+options
+
+tf.distribute.InputOptions used to control options on how this
+dataset is distributed.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.distribute.DistributedDataset.
+

+
+
experimental_assign_to_logical_device

+
+
+experimental_assign_to_logical_device(
+    tensor, logical_device_id
+)
+

+
+Adds annotation that `tensor` will be assigned to a logical device.
+
+This adds an annotation to `tensor` specifying that operations on `tensor` will
+be invoked on logical core device id `logical_device_id`. When model parallelism
+is used, the default behavior is that all ops are placed on zero-th logical
+device.
+
+```python
+
+# Initializing TPU system with 2 logical devices and 4 replicas.
+resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+tf.config.experimental_connect_to_cluster(resolver)
+topology = tf.tpu.experimental.initialize_tpu_system(resolver)
+device_assignment = tf.tpu.experimental.DeviceAssignment.build(
+    topology,
+    computation_shape=[1, 1, 1, 2],
+    num_replicas=4)
+strategy = tf.distribute.TPUStrategy(
+    resolver, experimental_device_assignment=device_assignment)
+iterator = iter(inputs)
+
+@tf.function()
+def step_fn(inputs):
+  output = tf.add(inputs, inputs)
+
+  # Add operation will be executed on logical device 0.
+  output = strategy.experimental_assign_to_logical_device(output, 0)
+  return output
+
+strategy.run(step_fn, args=(next(iterator),))
+```
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+tensor
+
+Input tensor to annotate.
+

+logical_device_id
+
+Id of the logical core to which the tensor will be
+assigned.
+

+
+
+
+ 
+
+
+
+
+
+
+
+
Raises

+ValueError
+
+The logical device id presented is not consistent with total
+number of partitions specified by the device assignment or the TPUStrategy
+is constructed with experimental_spmd_xla_partitioning=True.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+Annotated tensor with identical value as tensor.
+

+
+
experimental_distribute_dataset

+
+
+experimental_distribute_dataset(
+    dataset, options=None
+)
+

+
+Creates `tf.distribute.DistributedDataset` from `tf.data.Dataset`.
+
+The returned `tf.distribute.DistributedDataset` can be iterated over similar to
+regular datasets. NOTE: The user cannot add any more transformations to a
+`tf.distribute.DistributedDataset`. You can only create an iterator or examine
+the `tf.TypeSpec` of the data generated by it. See API docs of
+`tf.distribute.DistributedDataset` to learn more.
+
+The following is an example:
+
+```
+>>> global_batch_size = 2
+>>> # Passing the devices is optional.
+... strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
+>>> # Create a dataset
+... dataset = tf.data.Dataset.range(4).batch(global_batch_size)
+>>> # Distribute that dataset
+... dist_dataset = strategy.experimental_distribute_dataset(dataset)
+>>> @tf.function
+... def replica_fn(input):
+...   return input*2
+>>> result = []
+>>> # Iterate over the `tf.distribute.DistributedDataset`
+... for x in dist_dataset:
+...   # process dataset elements
+...   result.append(strategy.run(replica_fn, args=(x,)))
+>>> print(result)
+[PerReplica:{
+  0: ,
+  1: 
+}, PerReplica:{
+  0: ,
+  1: 
+}]
+```
+
+Three key actions happening under the hood of this method are batching,
+sharding, and prefetching.
+
+In the code snippet above, `dataset` is batched by `global_batch_size`, and
+calling `experimental_distribute_dataset` on it rebatches `dataset` to a new
+batch size that is equal to the global batch size divided by the number of
+replicas in sync. We iterate through it using a Pythonic for loop. `x` is a
+`tf.distribute.DistributedValues` containing data for all replicas, and each
+replica gets data of the new batch size. `tf.distribute.Strategy.run` will take
+care of feeding the right per-replica data in `x` to the right `replica_fn`
+executed on each replica.
+
+Sharding contains autosharding across multiple workers and within every worker.
+First, in multi-worker distributed training (i.e. when you use
+`tf.distribute.experimental.MultiWorkerMirroredStrategy` or
+`tf.distribute.TPUStrategy`), autosharding a dataset over a set of workers means
+that each worker is assigned a subset of the entire dataset (if the right
+`tf.data.experimental.AutoShardPolicy` is set). This is to ensure that at each
+step, a global batch size of non-overlapping dataset elements will be processed
+by each worker. Autosharding has a couple of different options that can be
+specified using `tf.data.experimental.DistributeOptions`. Then, sharding within
+each worker means the method will split the data among all the worker devices
+(if more than one a present). This will happen regardless of multi-worker
+autosharding.
+
+Note: for autosharding across multiple workers, the default mode is
+`tf.data.experimental.AutoShardPolicy.AUTO`. This mode will attempt to shard the
+input dataset by files if the dataset is being created out of reader datasets
+(e.g. `tf.data.TFRecordDataset`, `tf.data.TextLineDataset`, etc.) or otherwise
+shard the dataset by data, where each of the workers will read the entire
+dataset and only process the shard assigned to it. However, if you have less
+than one input file per worker, we suggest that you disable dataset autosharding
+across workers by setting the
+`tf.data.experimental.DistributeOptions.auto_shard_policy` to be
+`tf.data.experimental.AutoShardPolicy.OFF`.
+
+By default, this method adds a prefetch transformation at the end of the user
+provided `tf.data.Dataset` instance. The argument to the prefetch transformation
+which is `buffer_size` is equal to the number of replicas in sync.
+
+If the above batch splitting and dataset sharding logic is undesirable, please
+use `tf.distribute.Strategy.distribute_datasets_from_function` instead, which
+does not do any automatic batching or sharding for you.
+
+Note: If you are using TPUStrategy, the order in which the data is processed by
+the workers when using `tf.distribute.Strategy.experimental_distribute_dataset`
+or `tf.distribute.Strategy.distribute_datasets_from_function` is not guaranteed.
+This is typically required if you are using `tf.distribute` to scale prediction.
+You can however insert an index for each element in the batch and order outputs
+accordingly. Refer to
+[this snippet](https://www.tensorflow.org/tutorials/distribute/input#caveats)
+for an example of how to order outputs.
+
+Note: Stateful dataset transformations are currently not supported with
+`tf.distribute.experimental_distribute_dataset` or
+`tf.distribute.distribute_datasets_from_function`. Any stateful ops that the
+dataset may have are currently ignored. For example, if your dataset has a
+`map_fn` that uses `tf.random.uniform` to rotate an image, then you have a
+dataset graph that depends on state (i.e the random seed) on the local machine
+where the python process is being executed.
+
+For a tutorial on more usage and properties of this method, refer to the
+[tutorial on distributed input](https://www.tensorflow.org/tutorials/distribute/input#tfdistributestrategyexperimental_distribute_dataset).
+If you are interested in last partial batch handling, read
+[this section](https://www.tensorflow.org/tutorials/distribute/input#partial_batches).
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+dataset
+
+tf.data.Dataset that will be sharded across all replicas using
+the rules stated above.
+

+options
+
+tf.distribute.InputOptions used to control options on how this
+dataset is distributed.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.distribute.DistributedDataset.
+

+
+
experimental_distribute_values_from_function

+
+
+experimental_distribute_values_from_function(
+    value_fn
+)
+

+
+Generates `tf.distribute.DistributedValues` from `value_fn`.
+
+This function is to generate `tf.distribute.DistributedValues` to pass into
+`run`, `reduce`, or other methods that take distributed values when not using
+datasets.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+value_fn
+
+The function to run to generate values. It is called for
+each replica with tf.distribute.ValueContext as the sole argument. It
+must return a Tensor or a type that can be converted to a Tensor.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.distribute.DistributedValues containing a value for each replica.
+

+
+#### Example usage:
+
+1.  Return constant value per replica:
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> def value_fn(ctx):
+    ...   return tf.constant(1.)
+    >>> distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...        value_fn))
+    >>> local_result = strategy.experimental_local_results(
+    ...     distributed_values)
+    >>> local_result
+    (,
+    )
+    ```
+
+2.  Distribute values in array based on replica_id: {: value=2}
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> array_value = np.array([3., 2., 1.])
+    >>> def value_fn(ctx):
+    ...   return array_value[ctx.replica_id_in_sync_group]
+    >>> distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...         value_fn))
+    >>> local_result = strategy.experimental_local_results(
+    ...     distributed_values)
+    >>> local_result
+    (3.0, 2.0)
+    ```
+
+3.  Specify values using num_replicas_in_sync: {: value=3}
+
+    ```
+    >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+    >>> def value_fn(ctx):
+    ...   return ctx.num_replicas_in_sync
+    >>> distributed_values = (
+    ...     strategy.experimental_distribute_values_from_function(
+    ...         value_fn))
+    >>> local_result = strategy.experimental_local_results(
+    ...     distributed_values)
+    >>> local_result
+    (2, 2)
+    ```
+
+4.  Place values on devices and distribute: {: value=4}
+
+    ```
+    strategy = tf.distribute.TPUStrategy()
+    worker_devices = strategy.extended.worker_devices
+    multiple_values = []
+    for i in range(strategy.num_replicas_in_sync):
+      with tf.device(worker_devices[i]):
+        multiple_values.append(tf.constant(1.0))
+
+    def value_fn(ctx):
+      return multiple_values[ctx.replica_id_in_sync_group]
+
+    distributed_values = strategy.
+      experimental_distribute_values_from_function(
+      value_fn)
+    ```
+
+
experimental_local_results

+
+
+experimental_local_results(
+    value
+)
+

+
+Returns the list of all local per-replica values contained in `value`.
+
+Note: This only returns values on the worker initiated by this client. When
+using a `tf.distribute.Strategy` like
+`tf.distribute.experimental.MultiWorkerMirroredStrategy`, each worker will be
+its own client, and this function will only return values computed on that
+worker.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+value
+
+A value returned by experimental_run(), run(), or a variable
+created in scope`.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of values contained in value where ith element corresponds to
+ith replica. If value represents a single value, this returns
+(value,).
+

+
+
experimental_replicate_to_logical_devices

+
+
+experimental_replicate_to_logical_devices(
+    tensor
+)
+

+
+Adds annotation that `tensor` will be replicated to all logical devices.
+
+This adds an annotation to tensor `tensor` specifying that operations on
+`tensor` will be invoked on all logical devices.
+
+```python
+# Initializing TPU system with 2 logical devices and 4 replicas.
+resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+tf.config.experimental_connect_to_cluster(resolver)
+topology = tf.tpu.experimental.initialize_tpu_system(resolver)
+device_assignment = tf.tpu.experimental.DeviceAssignment.build(
+    topology,
+    computation_shape=[1, 1, 1, 2],
+    num_replicas=4)
+strategy = tf.distribute.TPUStrategy(
+    resolver, experimental_device_assignment=device_assignment)
+
+iterator = iter(inputs)
+
+@tf.function()
+def step_fn(inputs):
+  images, labels = inputs
+  images = strategy.experimental_split_to_logical_devices(
+    inputs, [1, 2, 4, 1])
+
+  # model() function will be executed on 8 logical devices with `inputs`
+  # split 2 * 4  ways.
+  output = model(inputs)
+
+  # For loss calculation, all logical devices share the same logits
+  # and labels.
+  labels = strategy.experimental_replicate_to_logical_devices(labels)
+  output = strategy.experimental_replicate_to_logical_devices(output)
+  loss = loss_fn(labels, output)
+
+  return loss
+
+strategy.run(step_fn, args=(next(iterator),))
+```
+
+Args: tensor: Input tensor to annotate.
+
+
+
+ 
+
+
+
+
+
+
+
Returns

+Annotated tensor with identical value as tensor.
+

+
+
experimental_split_to_logical_devices

+
+
+experimental_split_to_logical_devices(
+    tensor, partition_dimensions
+)
+

+
+Adds annotation that `tensor` will be split across logical devices.
+
+This adds an annotation to tensor `tensor` specifying that operations on
+`tensor` will be split among multiple logical devices. Tensor `tensor` will be
+split across dimensions specified by `partition_dimensions`. The dimensions of
+`tensor` must be divisible by corresponding value in `partition_dimensions`.
+
+For example, for system with 8 logical devices, if `tensor` is an image tensor
+with shape (batch_size, width, height, channel) and `partition_dimensions` is
+[1, 2, 4, 1], then `tensor` will be split 2 in width dimension and 4 way in
+height dimension and the split tensor values will be fed into 8 logical devices.
+
+```python
+# Initializing TPU system with 8 logical devices and 1 replica.
+resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+tf.config.experimental_connect_to_cluster(resolver)
+topology = tf.tpu.experimental.initialize_tpu_system(resolver)
+device_assignment = tf.tpu.experimental.DeviceAssignment.build(
+    topology,
+    computation_shape=[1, 2, 2, 2],
+    num_replicas=1)
+# Construct the TPUStrategy. Since we are going to split the image across
+# logical devices, here we set `experimental_spmd_xla_partitioning=True`
+# so that the partitioning can be compiled in SPMD mode, which usually
+# results in faster compilation and smaller HBM requirement if the size of
+# input and activation tensors are much bigger than that of the model
+# parameters. Note that this flag is suggested but not a hard requirement
+# for `experimental_split_to_logical_devices`.
+strategy = tf.distribute.TPUStrategy(
+    resolver, experimental_device_assignment=device_assignment,
+    experimental_spmd_xla_partitioning=True)
+
+iterator = iter(inputs)
+
+@tf.function()
+def step_fn(inputs):
+  inputs = strategy.experimental_split_to_logical_devices(
+    inputs, [1, 2, 4, 1])
+
+  # model() function will be executed on 8 logical devices with `inputs`
+  # split 2 * 4  ways.
+  output = model(inputs)
+  return output
+
+strategy.run(step_fn, args=(next(iterator),))
+```
+
+Args: tensor: Input tensor to annotate. partition_dimensions: An unnested list
+of integers with the size equal to rank of `tensor` specifying how `tensor` will
+be partitioned. The product of all elements in `partition_dimensions` must be
+equal to the total number of logical devices per replica.
+
+
+
+ 
+
+
+
+
+
+
+
+
Raises

+ValueError
+
+1) If the size of partition_dimensions does not equal to rank
+of tensor or 2) if product of elements of partition_dimensions does
+not match the number of logical devices per replica defined by the
+implementing DistributionStrategy's device specification or
+3) if a known size of tensor is not divisible by corresponding
+value in partition_dimensions.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+Annotated tensor with identical value as tensor.
+

+
+
gather

+
+
+gather(
+    value, axis
+)
+

+
+Gather `value` across replicas along `axis` to the current device.
+
+Given a `tf.distribute.DistributedValues` or `tf.Tensor`-like object `value`,
+this API gathers and concatenates `value` across replicas along the `axis`-th
+dimension. The result is copied to the "current" device, which would typically
+be the CPU of the worker on which the program is running. For
+`tf.distribute.TPUStrategy`, it is the first TPU host. For multi-client
+`tf.distribute.MultiWorkerMirroredStrategy`, this is the CPU of each worker.
+
+This API can only be called in the cross-replica context. For a counterpart in
+the replica context, see `tf.distribute.ReplicaContext.all_gather`.
+
+Note: For all strategies except `tf.distribute.TPUStrategy`, the input `value`
+on different replicas must have the same rank, and their shapes must be the same
+in all dimensions except the `axis`-th dimension. In other words, their shapes
+cannot be different in a dimension `d` where `d` does not equal to the `axis`
+argument. For example, given a `tf.distribute.DistributedValues` with component
+tensors of shape `(1, 2, 3)` and `(1, 3, 3)` on two replicas, you can call
+`gather(..., axis=1, ...)` on it, but not `gather(..., axis=0, ...)` or
+`gather(..., axis=2, ...)`. However, for `tf.distribute.TPUStrategy.gather`, all
+tensors must have exactly the same rank and same shape.
+
+Note: Given a `tf.distribute.DistributedValues` `value`, its component tensors
+must have a non-zero rank. Otherwise, consider using `tf.expand_dims` before
+gathering them.
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+>>> # A DistributedValues with component tensor of shape (2, 1) on each replica
+... distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
+>>> @tf.function
+... def run():
+...   return strategy.gather(distributed_values, axis=0)
+>>> run()
+
+```
+
+Consider the following example for more combinations:
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3"])
+>>> single_tensor = tf.reshape(tf.range(6), shape=(1,2,3))
+>>> distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(single_tensor))
+>>> @tf.function
+... def run(axis):
+...   return strategy.gather(distributed_values, axis=axis)
+>>> axis=0
+>>> run(axis)
+
+>>> axis=1
+>>> run(axis)
+
+>>> axis=2
+>>> run(axis)
+
+```
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+value
+
+a tf.distribute.DistributedValues instance, e.g. returned by
+Strategy.run, to be combined into a single tensor. It can also be a
+regular tensor when used with tf.distribute.OneDeviceStrategy or the
+default strategy. The tensors that constitute the DistributedValues
+can only be dense tensors with non-zero rank, NOT a tf.IndexedSlices.
+

+axis
+
+0-D int32 Tensor. Dimension along which to gather. Must be in the
+range [0, rank(value)).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A Tensor that's the concatenation of value across replicas along
+axis dimension.
+

+
+
reduce

+
+
+reduce(
+    reduce_op, value, axis
+)
+

+
+Reduce `value` across replicas and return result on current device.
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+>>> def step_fn():
+...   i = tf.distribute.get_replica_context().replica_id_in_sync_group
+...   return tf.identity(i)
+>>>
+>>> per_replica_result = strategy.run(step_fn)
+>>> total = strategy.reduce("SUM", per_replica_result, axis=None)
+>>> total
+
+```
+
+To see how this would look with multiple replicas, consider the same example
+with MirroredStrategy with 2 GPUs:
+
+```python
+strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
+def step_fn():
+  i = tf.distribute.get_replica_context().replica_id_in_sync_group
+  return tf.identity(i)
+
+per_replica_result = strategy.run(step_fn)
+# Check devices on which per replica result is:
+strategy.experimental_local_results(per_replica_result)[0].device
+# /job:localhost/replica:0/task:0/device:GPU:0
+strategy.experimental_local_results(per_replica_result)[1].device
+# /job:localhost/replica:0/task:0/device:GPU:1
+
+total = strategy.reduce("SUM", per_replica_result, axis=None)
+# Check device on which reduced result is:
+total.device
+# /job:localhost/replica:0/task:0/device:CPU:0
+
+```
+
+This API is typically used for aggregating the results returned from different
+replicas, for reporting etc. For example, loss computed from different replicas
+can be averaged using this API before printing.
+
+Note: The result is copied to the "current" device - which would typically be
+the CPU of the worker on which the program is running. For `TPUStrategy`, it is
+the first TPU host. For multi client `MultiWorkerMirroredStrategy`, this is CPU
+of each worker.
+
+There are a number of different tf.distribute APIs for reducing values across
+replicas: * `tf.distribute.ReplicaContext.all_reduce`: This differs from
+`Strategy.reduce` in that it is for replica context and does not copy the
+results to the host device. `all_reduce` should be typically used for reductions
+inside the training step such as gradients. *
+`tf.distribute.StrategyExtended.reduce_to` and
+`tf.distribute.StrategyExtended.batch_reduce_to`: These APIs are more advanced
+versions of `Strategy.reduce` as they allow customizing the destination of the
+result. They are also called in cross replica context.
+
+*What should axis be?*
+
+Given a per-replica value returned by `run`, say a per-example loss, the batch
+will be divided across all the replicas. This function allows you to aggregate
+across replicas and optionally also across batch elements by specifying the axis
+parameter accordingly.
+
+For example, if you have a global batch size of 8 and 2 replicas, values for
+examples `[0, 1, 2, 3]` will be on replica 0 and `[4, 5, 6, 7]` will be on
+replica 1. With `axis=None`, `reduce` will aggregate only across replicas,
+returning `[0+4, 1+5, 2+6, 3+7]`. This is useful when each replica is computing
+a scalar or some other value that doesn't have a "batch" dimension (like a
+gradient or loss). `strategy.reduce("sum", per_replica_result, axis=None)`
+
+Sometimes, you will want to aggregate across both the global batch *and* all
+replicas. You can get this behavior by specifying the batch dimension as the
+`axis`, typically `axis=0`. In this case it would return a scalar
+`0+1+2+3+4+5+6+7`. `strategy.reduce("sum", per_replica_result, axis=0)`
+
+If there is a last partial batch, you will need to specify an axis so that the
+resulting shape is consistent across replicas. So if the last batch has size 6
+and it is divided into [0, 1, 2, 3] and [4, 5], you would get a shape mismatch
+unless you specify `axis=0`. If you specify `tf.distribute.ReduceOp.MEAN`, using
+`axis=0` will use the correct denominator of 6. Contrast this with computing
+`reduce_mean` to get a scalar value on each replica and this function to average
+those means, which will weigh some values `1/8` and others `1/4`.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+reduce_op
+
+a tf.distribute.ReduceOp value specifying how values should
+be combined. Allows using string representation of the enum such as
+"SUM", "MEAN".
+

+value
+
+a tf.distribute.DistributedValues instance, e.g. returned by
+Strategy.run, to be combined into a single tensor. It can also be a
+regular tensor when used with OneDeviceStrategy or default strategy.
+

+axis
+
+specifies the dimension to reduce along within each
+replica's tensor. Should typically be set to the batch dimension, or
+None to only reduce across replicas (e.g. if the tensor has no batch
+dimension).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A Tensor.
+

+
+
run

+
+
+run(
+    fn, args=(), kwargs=None, options=None
+)
+

+
+Run the computation defined by `fn` on each TPU replica.
+
+Executes ops specified by `fn` on each replica. If `args` or `kwargs` have
+`tf.distribute.DistributedValues`, such as those produced by a
+`tf.distribute.DistributedDataset` from
+`tf.distribute.Strategy.experimental_distribute_dataset` or
+`tf.distribute.Strategy.distribute_datasets_from_function`, when `fn` is
+executed on a particular replica, it will be executed with the component of
+`tf.distribute.DistributedValues` that correspond to that replica.
+
+`fn` may call `tf.distribute.get_replica_context()` to access members such as
+`all_reduce`.
+
+All arguments in `args` or `kwargs` should either be nest of tensors or
+`tf.distribute.DistributedValues` containing tensors or composite tensors.
+
+#### Example usage:
+
+```
+>>> resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+>>> tf.config.experimental_connect_to_cluster(resolver)
+>>> tf.tpu.experimental.initialize_tpu_system(resolver)
+>>> strategy = tf.distribute.TPUStrategy(resolver)
+>>> @tf.function
+... def run():
+...   def value_fn(value_context):
+...     return value_context.num_replicas_in_sync
+...   distributed_values = (
+...       strategy.experimental_distribute_values_from_function(value_fn))
+...   def replica_fn(input):
+...     return input * 2
+...   return strategy.run(replica_fn, args=(distributed_values,))
+>>> result = run()
+```
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+fn
+
+The function to run. The output must be a tf.nest of Tensors.
+

+args
+
+(Optional) Positional arguments to fn.
+

+kwargs
+
+(Optional) Keyword arguments to fn.
+

+options
+
+(Optional) An instance of tf.distribute.RunOptions specifying
+the options to run fn.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+Merged return value of fn across replicas. The structure of the return
+value is the same as the return value from fn. Each element in the
+structure can either be tf.distribute.DistributedValues, Tensor
+objects, or Tensors (for example, if running on a single replica).
+

+
+
scope

+
+
+scope()
+

+
+Context manager to make the strategy current and distribute variables.
+
+This method returns a context manager, and is used as follows:
+
+```
+>>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+>>> # Variable created inside scope:
+>>> with strategy.scope():
+...   mirrored_variable = tf.Variable(1.)
+>>> mirrored_variable
+MirroredVariable:{
+  0: ,
+  1: 
+}
+>>> # Variable created outside scope:
+>>> regular_variable = tf.Variable(1.)
+>>> regular_variable
+
+```
+
+*What happens when Strategy.scope is entered?*
+
+*   `strategy` is installed in the global context as the "current" strategy.
+    Inside this scope, `tf.distribute.get_strategy()` will now return this
+    strategy. Outside this scope, it returns the default no-op strategy.
+*   Entering the scope also enters the "cross-replica context". See
+    `tf.distribute.StrategyExtended` for an explanation on cross-replica and
+    replica contexts.
+*   Variable creation inside `scope` is intercepted by the strategy. Each
+    strategy defines how it wants to affect the variable creation. Sync
+    strategies like `MirroredStrategy`, `TPUStrategy` and
+    `MultiWorkerMiroredStrategy` create variables replicated on each replica,
+    whereas `ParameterServerStrategy` creates variables on the parameter
+    servers. This is done using a custom `tf.variable_creator_scope`.
+*   In some strategies, a default device scope may also be entered: in
+    `MultiWorkerMiroredStrategy`, a default device scope of "/CPU:0" is entered
+    on each worker.
+
+Note: Entering a scope does not automatically distribute a computation, except
+in the case of high level training framework like keras `model.fit`. If you're
+not using `model.fit`, you need to use `strategy.run` API to explicitly
+distribute that computation. See an example in the
+[custom training loop tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training).
+
+*What should be in scope and what should be outside?*
+
+There are a number of requirements on what needs to happen inside the scope.
+However, in places where we have information about which strategy is in use, we
+often enter the scope for the user, so they don't have to do it explicitly (i.e.
+calling those either inside or outside the scope is OK).
+
+*   Anything that creates variables that should be distributed variables must be
+    called in a `strategy.scope`. This can be accomplished either by directly
+    calling the variable creating function within the scope context, or by
+    relying on another API like `strategy.run` or `keras.Model.fit` to
+    automatically enter it for you. Any variable that is created outside scope
+    will not be distributed and may have performance implications. Some common
+    objects that create variables in TF are Models, Optimizers, Metrics. Such
+    objects should always be initialized in the scope, and any functions that
+    may lazily create variables (e.g., `Model.__call__()`, tracing a
+    `tf.function`, etc.) should similarly be called within scope. Another source
+    of variable creation can be a checkpoint restore - when variables are
+    created lazily. Note that any variable created inside a strategy captures
+    the strategy information. So reading and writing to these variables outside
+    the `strategy.scope` can also work seamlessly, without the user having to
+    enter the scope.
+*   Some strategy APIs (such as `strategy.run` and `strategy.reduce`) which
+    require to be in a strategy's scope, enter the scope automatically, which
+    means when using those APIs you don't need to explicitly enter the scope
+    yourself.
+*   When a `tf.keras.Model` is created inside a `strategy.scope`, the Model
+    object captures the scope information. When high level training framework
+    methods such as `model.compile`, `model.fit`, etc. are then called, the
+    captured scope will be automatically entered, and the associated strategy
+    will be used to distribute the training etc. See a detailed example in
+    [distributed keras tutorial](https://www.tensorflow.org/tutorials/distribute/keras).
+    WARNING: Simply calling `model(..)` does not automatically enter the
+    captured scope -- only high level training framework APIs support this
+    behavior: `model.compile`, `model.fit`, `model.evaluate`, `model.predict`
+    and `model.save` can all be called inside or outside the scope.
+*   The following can be either inside or outside the scope:
+    *   Creating the input datasets
+    *   Defining `tf.function`s that represent your training step
+    *   Saving APIs such as `tf.saved_model.save`. Loading creates variables, so
+        that should go inside the scope if you want to train the model in a
+        distributed way.
+    *   Checkpoint saving. As mentioned above - `checkpoint.restore` may
+        sometimes need to be inside scope if it creates variables.
+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A context manager.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/Task.md b/tensorflow_gnn/docs/api_docs/python/runner/Task.md
new file mode 100644
index 00000000..770bcb48
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/Task.md
@@ -0,0 +1,179 @@
+# runner.Task
+
+
+
+
+ View source
+on GitHub 
+
+Defines a learning objective for a GNN.
+
+
+
+A `Task` represents a learning objective for a GNN model and defines all the
+non-GNN pieces around the base GNN. Specifically:
+
+1) `preprocess` is expected to return a `GraphTensor` (or `GraphTensor`s) and a
+`Field` where (a) the base GNN's output for each `GraphTensor` is passed to
+`predict` and (b) the `Field` is used as the training label (for supervised
+tasks); 2) `predict` is expected to (a) take the base GNN's output for each
+`GraphTensor` returned by `preprocess` and (b) return a tensor with the model's
+prediction for this task; 3) `losses` is expected to return callables
+(`tf.Tensor`, `tf.Tensor`) -> `tf.Tensor` that accept (`y_true`, `y_pred`) where
+`y_true` is produced by some dataset and `y_pred` is the model's prediction from
+(2); 4) `metrics` is expected to return callables (`tf.Tensor`, `tf.Tensor`) ->
+`tf.Tensor` that accept (`y_true`, `y_pred`) where `y_true` is produced by some
+dataset and `y_pred` is the model's prediction from (2).
+
+`Task` can emit multiple outputs in `predict`: in that case we require that (a)
+it is a mapping, (b) outputs of `losses` and `metrics` are also mappings with
+matching keys, and (c) there is exactly one loss per key (there may be a
+sequence of metrics per key). This is done to prevent accidental dropping of
+losses (see b/291874188).
+
+No constraints are made on the `predict` method; e.g.: it may append a head with
+learnable weights or it may perform tensor computations only. (The entire `Task`
+coordinates what that means with respect to dataset—via `preprocess`—,
+modeling—via `predict`— and optimization—via `losses`.)
+
+`Task`s are applied in the scope of a training invocation: they are subject to
+the executing context of the `Trainer` and should, when needed, override it
+(e.g., a global policy, like `tf.keras.mixed_precision.global_policy()` and its
+implications over logit and activation layers).
+
+## Methods
+
+
losses

+
+View
+source
+
+
+@abc.abstractmethod
+losses() -> Losses
+

+
+Returns arbitrary task specific losses.
+
+
metrics

+
+View
+source
+
+
+@abc.abstractmethod
+metrics() -> Metrics
+

+
+Returns arbitrary task specific metrics.
+
+
predict

+
+View
+source
+
+
+@abc.abstractmethod
+predict(
+    *args
+) -> Predictions
+

+
+Produces prediction outputs for the learning objective.
+
+Overall model composition* makes use of the Keras Functional API
+(https://www.tensorflow.org/guide/keras/functional) to map symbolic Keras
+`GraphTensor` inputs to symbolic Keras `Field` outputs. Outputs must match the
+structure (one or mapping) of labels from `preprocess`.
+
+*) `outputs = predict(GNN(inputs))` where `inputs` are those `GraphTensor`
+returned by `preprocess(...)`, `GNN` is the base GNN, `predict` is this method
+and `outputs` are the prediction outputs for the learning objective.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+*args
+
+The symbolic Keras GraphTensor inputs(s). These inputs correspond
+(in sequence) to the base GNN output of each GraphTensor returned by
+preprocess(...).
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The model's prediction output for this task.
+

+
+
preprocess

+
+View
+source
+
+
+@abc.abstractmethod
+preprocess(
+    inputs: GraphTensor
+) -> tuple[OneOrSequenceOf[GraphTensor], OneOrMappingOf[Field]]
+

+
+Preprocesses a scalar (after `merge_batch_to_components`) `GraphTensor`.
+
+This function uses the Keras functional API to define non-trainable
+transformations of the symbolic input `GraphTensor`, which get executed during
+dataset preprocessing in a `tf.data.Dataset.map(...)` operation. It has two
+responsibilities:
+
+1.  Splitting the training label out of the input for training. It must be
+    returned as a separate tensor or mapping of tensors.
+2.  Optionally, transforming input features. Some advanced modeling techniques
+    require running the same base GNN on multiple different transformations, so
+    this function may return a single `GraphTensor` or a non-empty sequence of
+    `GraphTensors`. The corresponding base GNN output for each `GraphTensor` is
+    provided to the `predict(...)` method.
+
+
+
+ 
+
+
+
+
+
+
+
+
Args

+inputs
+
+A symbolic Keras GraphTensor for processing.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tuple of processed GraphTensor(s) and a (one or mapping of) Field to
+be used as labels.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/TightPadding.md b/tensorflow_gnn/docs/api_docs/python/runner/TightPadding.md
new file mode 100644
index 00000000..d859478b
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/TightPadding.md
@@ -0,0 +1,47 @@
+# runner.TightPadding
+
+
+
+
+ View source
+on GitHub 
+
+Calculates tight `SizeConstraints` for `GraphTensor` padding.
+
+Inherits From: [`GraphTensorPadding`](../runner/GraphTensorPadding.md)
+
+
+runner.TightPadding(
+    gtspec: tfgnn.GraphTensorSpec,
+    dataset_provider: runner.DatasetProvider,
+    min_nodes_per_component: Optional[Mapping[str, int]] = None
+)
+

+
+
+
+See: `tfgnn.find_tight_size_constraints.`
+
+## Methods
+
+
get_filter_fn

+
+View
+source
+
+
+get_filter_fn(
+    size_constraints: SizeConstraints
+) -> Callable[..., bool]
+

+
+
get_size_constraints

+
+View
+source
+
+
+get_size_constraints(
+    target_batch_size: int
+) -> SizeConstraints
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/Trainer.md b/tensorflow_gnn/docs/api_docs/python/runner/Trainer.md
new file mode 100644
index 00000000..1b1db5f8
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/Trainer.md
@@ -0,0 +1,100 @@
+# runner.Trainer
+
+
+
+
+ View source
+on GitHub 
+
+A class for training and validation of a Keras model.
+
+
+
+
+ 
+
+
+
+
+
+
Attributes
  model_dir  
+
+ 
  strategy  
+
+

+
+## Methods
+
+
train

+
+View
+source
+
+
+@abc.abstractmethod
+train(
+    model_fn: Callable[[], tf.keras.Model],
+    train_ds_provider: DatasetProvider,
+    *,
+    epochs: int = 1,
+    valid_ds_provider: Optional[DatasetProvider] = None
+) -> tf.keras.Model
+

+
+Trains a `tf.keras.Model` with optional validation.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+model_fn
+
+Returns a tf.keras.Model for use in training and validation.
+

+train_ds_provider
+
+A DatasetProvider for training. The items of the
+tf.data.Dataset are pairs (graph_tensor, label) that represent one
+batch of per-replica training inputs after
+GraphTensor.merge_batch_to_components() has been applied.
+

+epochs
+
+The epochs to train.
+

+valid_ds_provider
+
+A DatasetProvider for validation. The items of the
+tf.data.Dataset are pairs (graph_tensor, label) that represent one
+batch of per-replica training inputs after
+GraphTensor.merge_batch_to_components() has been applied.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A trained tf.keras.Model.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/all_symbols.md b/tensorflow_gnn/docs/api_docs/python/runner/all_symbols.md
new file mode 100644
index 00000000..b8281392
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/all_symbols.md
@@ -0,0 +1,62 @@
+# All symbols in TensorFlow GNN Runner
+
+
+
+## Primary symbols
+
+*   runner
+*   runner.ContextLabelFn
+*   runner.DatasetProvider
+*   runner.DotProductLinkPrediction
+*   runner.FitOrSkipPadding
+*   runner.GraphBinaryClassification
+*   runner.GraphMeanAbsoluteError
+*   runner.GraphMeanAbsolutePercentageError
+*   runner.GraphMeanSquaredError
+*   runner.GraphMeanSquaredLogScaledError
+*   runner.GraphMeanSquaredLogarithmicError
+*   runner.GraphMulticlassClassification
+*   runner.GraphTensorPadding
+*   runner.GraphTensorProcessorFn
+*   runner.HadamardProductLinkPrediction
+*   runner.IntegratedGradientsExporter
+*   runner.KerasModelExporter
+*   runner.KerasTrainer
+*   runner.KerasTrainerCheckpointOptions
+*   runner.KerasTrainerOptions
+*   runner.Loss
+*   runner.Losses
+*   runner.Metric
+*   runner.Metrics
+*   runner.ModelExporter
+*   runner.NodeBinaryClassification
+*   runner.NodeMulticlassClassification
+*   runner.ParameterServerStrategy
+*   runner.PassthruDatasetProvider
+*   runner.PassthruSampleDatasetsProvider
+*   runner.Predictions
+*   runner.RootNodeBinaryClassification
+*   runner.RootNodeLabelFn
+*   runner.RootNodeMeanAbsoluteError
+*   runner.RootNodeMeanAbsoluteLogarithmicError
+*   runner.RootNodeMeanAbsolutePercentageError
+*   runner.RootNodeMeanSquaredError
+*   runner.RootNodeMeanSquaredLogScaledError
+*   runner.RootNodeMeanSquaredLogarithmicError
+*   runner.RootNodeMulticlassClassification
+*   runner.RunResult
+*   runner.SampleTFRecordDatasetsProvider
+*   runner.SimpleDatasetProvider
+*   runner.SimpleSampleDatasetsProvider
+*   runner.SubmoduleExporter
+*   runner.TFDataServiceConfig
+*   runner.TFRecordDatasetProvider
+*   runner.TPUStrategy
+*   runner.Task
+*   runner.TightPadding
+*   runner.Trainer
+*   runner.export_model
+*   runner.incrementing_model_dir
+*   runner.integrated_gradients
+*   runner.one_node_per_component
+*   runner.run
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/export_model.md b/tensorflow_gnn/docs/api_docs/python/runner/export_model.md
new file mode 100644
index 00000000..22c27634
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/export_model.md
@@ -0,0 +1,77 @@
+# runner.export_model
+
+
+
+
+ View source
+on GitHub 
+
+Exports a Keras model without traces s.t. it is loadable without TF-GNN.
+
+
+runner.export_model(
+    model: tf.keras.Model,
+    export_dir: str,
+    *,
+    output_names: Optional[Any] = None,
+    options: Optional[tf.saved_model.SaveOptions] = None,
+    use_legacy_model_save: Optional[bool] = None
+) -> None
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+model
+
+Keras model instance to be saved.
+

+export_dir
+
+Path where to save the model.
+

+output_names
+
+Optionally, a nest of str values or None with the same
+structure as the outputs of model. A non-None value is used as that
+output's key in the SavedModel signature. By default, an output gets
+the name of the final Keras layer creating it as its key (matching the
+behavior of legacy Model.save(save_format="tf")).
+

+options
+
+An optional tf.saved_model.SaveOptions argument.
+

+use_legacy_model_save
+
+Optional; most users can leave it unset to get a
+useful default for export to inference. If set to True, forces the use
+of Model.save(), which exports a SavedModel suitable for inference and
+potentially also for reloading as a Keras model (depending on its Layers).
+If set to False, forces the use of tf.keras.export.ExportArchive,
+which is usable as of TensorFlow 2.13 and is advertised as the more
+streamlined way of exporting to SavedModel for inference only. Currently,
+None behaves like True, but the long-term plan is to migrate towards
+False.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/incrementing_model_dir.md b/tensorflow_gnn/docs/api_docs/python/runner/incrementing_model_dir.md
new file mode 100644
index 00000000..701754de
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/incrementing_model_dir.md
@@ -0,0 +1,52 @@
+# runner.incrementing_model_dir
+
+
+
+
+ View source
+on GitHub 
+
+Create, given some `dirname`, an incrementing model directory.
+
+
+runner.incrementing_model_dir(
+    dirname: str, start: int = 0
+) -> str
+

+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Args

+dirname
+
+The base directory name.
+

+start
+
+The starting integer.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A model directory dirname/n where 'n' is the maximum integer in dirname.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/integrated_gradients.md b/tensorflow_gnn/docs/api_docs/python/runner/integrated_gradients.md
new file mode 100644
index 00000000..7b2c69ef
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/integrated_gradients.md
@@ -0,0 +1,95 @@
+# runner.integrated_gradients
+
+
+
+
+ View source
+on GitHub 
+
+Integrated gradients.
+
+
+runner.integrated_gradients(
+    preprocess_model: tf.keras.Model,
+    model: tf.keras.Model,
+    *,
+    output_name: Optional[str] = None,
+    random_counterfactual: bool,
+    steps: int,
+    seed: Optional[int] = None
+) -> tf.types.experimental.ConcreteFunction
+

+
+
+
+This `tf.function` computes integrated gradients over a `tfgnn.GraphTensor.` The
+`tf.function` will be persisted in the ultimate saved model for subsequent
+attribution.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+preprocess_model
+
+A tf.keras.Model for preprocessing. This model is
+expected to return a tuple (GraphTensor, Tensor) where the
+GraphTensor is used to invoke the below model and the tensor is used
+used for any loss computation. (Via model.compiled_loss.)
+

+model
+
+A tf.keras.Model for integrated gradients.
+

+output_name
+
+The output Tensor name. If unset, the tensor will be named
+by Keras defaults.
+

+random_counterfactual
+
+Whether to use a random uniform counterfactual.
+

+steps
+
+The number of interpolations of the Riemann sum approximation.
+

+seed
+
+An option random seed.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A tf.function with the integrated gradients as output.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/runner/one_node_per_component.md b/tensorflow_gnn/docs/api_docs/python/runner/one_node_per_component.md
new file mode 100644
index 00000000..a260d8b0
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/one_node_per_component.md
@@ -0,0 +1,17 @@
+# runner.one_node_per_component
+
+
+
+
+ View source
+on GitHub 
+
+Returns a `Mapping` `node_set_name: 1` for every node set in `gtspec`.
+
+
+runner.one_node_per_component(
+    gtspec: tfgnn.GraphTensorSpec
+) -> Mapping[str, int]
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/runner/run.md b/tensorflow_gnn/docs/api_docs/python/runner/run.md
new file mode 100644
index 00000000..7c234650
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/runner/run.md
@@ -0,0 +1,275 @@
+# runner.run
+
+
+
+
+ View source
+on GitHub 
+
+Runs training (and validation) of a model on task(s) with the given data.
+
+
+runner.run(
+    *,
+    train_ds_provider: DatasetProvider,
+    model_fn: Callable[[GraphTensorSpec], tf.keras.Model],
+    optimizer_fn: Callable[[], tf.keras.optimizers.Optimizer],
+    trainer: Trainer,
+    task: OneOrMappingOf[Task],
+    loss_weights: Optional[Mapping[str, float]] = None,
+    gtspec: GraphTensorSpec,
+    global_batch_size: int,
+    epochs: int = 1,
+    drop_remainder: bool = False,
+    export_dirs: Optional[Sequence[str]] = None,
+    model_exporters: Optional[Sequence[ModelExporter]] = None,
+    feature_processors: Optional[Sequence[GraphTensorProcessorFn]] = None,
+    valid_ds_provider: Optional[DatasetProvider] = None,
+    train_padding: Optional[GraphTensorPadding] = None,
+    valid_padding: Optional[GraphTensorPadding] = None,
+    tf_data_service_config: Optional[TFDataServiceConfig] = None,
+    steps_per_execution: Optional[int] = None,
+    run_eagerly: bool = False
+)
+

+
+
+
+This includes preprocessing the input data, appending any suitable head(s), and
+running training (and validation) with the requested distribution strategy.
+
+The input data is processed in multiple stages, starting from the contents of
+the datasets provided by `train_ds_provider` and `valid_ds_provider`:
+
+1.  Input examples are batched.
+2.  If necessary, input batches are parsed as `GraphTensor` values and merged
+    into components (see: `GraphTensor.merge_batch_to_components`).
+3.  If set, `train_padding` and `valid_padding`, resp., are applied.
+4.  The given `feature_processors` are applied in order for all non-trainable
+    feature transformations on CPU (as part of `tf.data.Dataset.map(...)`).
+5.  The
+    Task.preprocess(...)
+    method is applied to extract training targets (for supervised learning, that
+    means: labels) and optionally transform the value of the preprocessed
+    `GraphTensor` into a model input (or multiple model inputs for tasks like
+    self-supervised contrastive losses).
+6.  If the resulting `GraphTensor`s have any auxiliary pieces (as indicated by
+    `tfgnn.get_aux_type_prefix(...)`): all features (typically: labels) are
+    removed from those graph pieces.
+
+The base GNN (as built by `model_fn`) is run on all results from step (6).
+Task.predict(...) is called
+on the model outputs that correspond to the one or more graphs requested in step
+(5) by
+Task.preprocess(...).
+
+Trainable transformations of inputs (notably lookups in trainable embedding
+tables) are required to happen inside `model_fn`.
+
+For supervised learning, training labels enter the pipeline as features on the
+`GraphTensor` that undergo the `feature_processors` (shared by all `Task`s) and
+are read out of the `GraphTensor` by
+Task.preprocess(...).
+
+Users are strongly encouraged to take one of the following two approaches to
+prevent the leakage of label information into the training:
+
+*   Store labels on the auxiliary `"_readout"` node set and let
+    Task.preprocess(...)
+    read them from there. (For library-supplied `Task`s, that means initializing
+    with `label_feature_name="..."`.) If that is not already true for the input
+    datasets, the label feature can be moved there by one of the
+    `feature_processors`, using `tfgnn.structured_readout_into_feature(...)` or
+    a similar helper function.
+*   For single-Task training only: Let
+    Task.preprocess()
+    return modified `GraphTensor`s that no longer contain the separately
+    returned labels. (Library-supplied Tasks delegate this to the
+    `label_fn="..."` passed in initialization.)
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+train_ds_provider
+
+A DatasetProvider for training. The tf.data.Dataset
+is not batched and contains scalar GraphTensor values conforming to
+gtspec, possibly serialized as a tf.train.Example proto.
+

+model_fn
+
+Returns the base GNN tf.keras.Model for use in training and
+validation.
+

+optimizer_fn
+
+Returns a tf.keras.optimizers.Optimizer for use in training.
+

+trainer
+
+A Trainer.
+

+task
+
+A Task for single-Task training or a Mapping[str, Task] for
+multi-Task training. In multi-Task training, Task.preprocess(...)
+must return GraphTensors with the same spec as its inputs, only the
+values may change (so that there remains a single spec for model_fn).
+

+loss_weights
+
+An optional Mapping[str, float] for multi-Task training. If
+given, this structure must match (with tf.nest.assert_same_structure)
+the structure of task. The mapping contains, for each task, a scalar
+coefficient to weight the loss contributions of that task.
+

+gtspec
+
+A GraphTensorSpec matching the elements of train and valid
+datasets. If train or valid contain tf.string elements, this
+GraphTensorSpec is used for parsing; otherwise, train or valid are
+expected to contain GraphTensor elements whose relaxed spec matches
+gtspec.
+

+global_batch_size
+
+The tf.data.Dataset global batch size for both training
+and validation.
+

+epochs
+
+The epochs to train.
+

+drop_remainder
+
+Whether to drop a tf.data.Dataset remainder at batching.
+

+export_dirs
+
+Optional directories for exports (SavedModels); if unset,
+default behavior is os.path.join(model_dir, "export").
+

+model_exporters
+
+Zero or more ModelExporter for exporting (SavedModels) to
+export_dirs. If unset, default behavior is [KerasModelExporter()].
+

+feature_processors
+
+A sequence of callables for feature processing with the
+Keras functional API. Each callable must accept and return a symbolic
+scalar GraphTensor. The callables are composed in order and may change
+the GraphTensorSpec (e.g., add/remove features). The resulting Keras
+model is executed on CPU as part of a tf.data.Dataset.map operation.
+

+valid_ds_provider
+
+A DatasetProvider for validation. The tf.data.Dataset
+is not batched and contains scalar GraphTensor values conforming to
+gtspec, possibly serialized as a tf.train.Example proto.
+

+train_padding
+
+GraphTensor padding for training. Required if training on
+TPU.
+

+valid_padding
+
+GraphTensor padding for validation. Required if training on
+TPU.
+

+tf_data_service_config
+
+tf.data service speeds-up tf.data input pipeline
+runtime reducing input bottlenecks for model training. Particularly for
+training on accelerators consider enabling it. For more info please see:
+https://www.tensorflow.org/api_docs/python/tf/data/experimental/service.
+

+steps_per_execution
+
+The number of batches to run during each training
+iteration. If not set, for TPU strategy default to 100 and to None
+otherwise.
+

+run_eagerly
+
+Whether to compile the model in eager mode, primarily for
+debugging purposes. Note that the symbolic model will still be run twice,
+so if you use a breakpoint() you will have to Continue twice before you
+are in a real eager execution.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+A RunResult object containing models and information about this run.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn.md b/tensorflow_gnn/docs/api_docs/python/tfgnn.md
index 43d214c5..2d03c918 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn.md
@@ -1,32 +1,20 @@
 # Module: tfgnn
 
-[TOC]
-
 
 
-
-
-
-
-
+
+ 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
-
-The various data types provided by the GNN library have corresponding schemas
-similar to `tf.TensorSpec`. For example, a `FieldSpec` describes an instance of
-`Field`, and a `GraphTensorSpec` describes an instance of `GraphTensor`.
+```
+import tensorflow_gnn as tfgnn
+```
 
 ## Modules
 
@@ -35,43 +23,63 @@ of the public interface of TensorFlow GNN.
 
 [`keras`](./tfgnn/keras.md) module: The tfgnn.keras package.
 
+[`proto`](./tfgnn/proto.md) module: The protocol message (protobuf) types
+defined by TensorFlow GNN.
+
 [`sampler`](./tfgnn/sampler.md) module: Public interface for GNN Sampler.
 
 ## Classes
 
-[`class Adjacency`](./tfgnn/Adjacency.md): Stores how edges connect pairs of nodes from source and target node sets.
+[`class Adjacency`](./tfgnn/Adjacency.md): Stores how edges connect pairs of
+nodes from source and target node sets.
 
-[`class AdjacencySpec`](./tfgnn/AdjacencySpec.md): A type spec for tfgnn.Adjacency.
+[`class AdjacencySpec`](./tfgnn/AdjacencySpec.md): A type spec for
+tfgnn.Adjacency.
 
-[`class Context`](./tfgnn/Context.md): A composite tensor for graph context features.
+[`class Context`](./tfgnn/Context.md): A composite tensor for graph context
+features.
 
-[`class ContextSpec`](./tfgnn/ContextSpec.md): A type spec for tfgnn.Context.
+[`class ContextSpec`](./tfgnn/ContextSpec.md): A type spec for
+tfgnn.Context.
 
-[`class EdgeSet`](./tfgnn/EdgeSet.md): A composite tensor for edge set features, size and adjacency information.
+[`class EdgeSet`](./tfgnn/EdgeSet.md): A composite tensor for edge set features,
+size and adjacency information.
 
-[`class EdgeSetSpec`](./tfgnn/EdgeSetSpec.md): A type spec for tfgnn.EdgeSet.
+[`class EdgeSetSpec`](./tfgnn/EdgeSetSpec.md): A type spec for
+tfgnn.EdgeSet.
 
-[`class Feature`](./tfgnn/Feature.md): A schema for a single feature.
+[`class Feature`](./tfgnn/proto/Feature.md): The schema entry for a single
+feature.
 
-[`class FeatureDefaultValues`](./tfgnn/FeatureDefaultValues.md): Default values for graph context, node sets and edge sets features.
+[`class FeatureDefaultValues`](./tfgnn/FeatureDefaultValues.md): Default values
+for graph context, node sets and edge sets features.
 
-[`class GraphSchema`](./tfgnn/GraphSchema.md): A schema definition for graphs.
+[`class GraphSchema`](./tfgnn/proto/GraphSchema.md): The top-level container for
+the schema of a graph dataset.
 
-[`class GraphTensor`](./tfgnn/GraphTensor.md): A composite tensor for heterogeneous directed graphs with features.
+[`class GraphTensor`](./tfgnn/GraphTensor.md): A composite tensor for
+heterogeneous directed graphs with features.
 
-[`class GraphTensorSpec`](./tfgnn/GraphTensorSpec.md): A type spec for tfgnn.GraphTensor.
+[`class GraphTensorSpec`](./tfgnn/GraphTensorSpec.md): A type spec for
+tfgnn.GraphTensor.
 
-[`class HyperAdjacency`](./tfgnn/HyperAdjacency.md): Stores how (hyper-)edges connect tuples of nodes from incident node sets.
+[`class HyperAdjacency`](./tfgnn/HyperAdjacency.md): Stores how (hyper-)edges
+connect tuples of nodes from incident node sets.
 
-[`class HyperAdjacencySpec`](./tfgnn/HyperAdjacencySpec.md): A type spec for tfgnn.HyperAdjacency.
+[`class HyperAdjacencySpec`](./tfgnn/HyperAdjacencySpec.md): A type spec for
+tfgnn.HyperAdjacency.
 
-[`class NodeSet`](./tfgnn/NodeSet.md): A composite tensor for node set features plus size information.
+[`class NodeSet`](./tfgnn/NodeSet.md): A composite tensor for node set features
+plus size information.
 
-[`class NodeSetSpec`](./tfgnn/NodeSetSpec.md): A type spec for tfgnn.NodeSet.
+[`class NodeSetSpec`](./tfgnn/NodeSetSpec.md): A type spec for
+tfgnn.NodeSet.
 
-[`class SizeConstraints`](./tfgnn/SizeConstraints.md): Constraints on the number of entities in the graph.
+[`class SizeConstraints`](./tfgnn/SizeConstraints.md): Constraints on the number
+of entities in the graph.
 
-[`class ValidationError`](./tfgnn/ValidationError.md): A schema validation error.
+[`class ValidationError`](./tfgnn/ValidationError.md): A schema validation
+error.
 
 ## Functions
 
@@ -79,23 +87,29 @@ of the public interface of TensorFlow GNN.
 Adds a readout structure equivalent to
 tfgnn.gather_first_node().
 
-[`add_self_loops(...)`](./tfgnn/add_self_loops.md): Adds self-loops for edge
-with name `edge_set_name` EVEN if already exist.
+[`add_self_loops(...)`](./tfgnn/add_self_loops.md): Adds self-loops for
+`edge_set_name` EVEN if they already exist.
 
-[`assert_constraints(...)`](./tfgnn/assert_constraints.md): Validate the shape constaints of a graph's features at runtime.
+[`assert_constraints(...)`](./tfgnn/assert_constraints.md): Validate the shape
+constaints of a graph's features at runtime.
 
-[`assert_satisfies_size_constraints(...)`](./tfgnn/assert_satisfies_size_constraints.md): Raises InvalidArgumentError if graph_tensor exceeds size_constraints.
+[`assert_satisfies_size_constraints(...)`](./tfgnn/assert_satisfies_size_constraints.md):
+Raises InvalidArgumentError if graph_tensor exceeds size_constraints.
 
-[`assert_satisfies_total_sizes(...)`](./tfgnn/assert_satisfies_size_constraints.md): Raises InvalidArgumentError if graph_tensor exceeds size_constraints.
+[`assert_satisfies_total_sizes(...)`](./tfgnn/assert_satisfies_size_constraints.md):
+Raises InvalidArgumentError if graph_tensor exceeds size_constraints.
 
 [`broadcast(...)`](./tfgnn/broadcast.md): Broadcasts values from nodes to edges,
 or from context to nodes or edges.
 
-[`broadcast_context_to_edges(...)`](./tfgnn/broadcast_context_to_edges.md): Broadcasts a context value to the `edge_set` edges.
+[`broadcast_context_to_edges(...)`](./tfgnn/broadcast_context_to_edges.md):
+Broadcasts a context value to the `edge_set` edges.
 
-[`broadcast_context_to_nodes(...)`](./tfgnn/broadcast_context_to_nodes.md): Broadcasts a context value to the `node_set` nodes.
+[`broadcast_context_to_nodes(...)`](./tfgnn/broadcast_context_to_nodes.md):
+Broadcasts a context value to the `node_set` nodes.
 
-[`broadcast_node_to_edges(...)`](./tfgnn/broadcast_node_to_edges.md): Broadcasts values from nodes to incident edges.
+[`broadcast_node_to_edges(...)`](./tfgnn/broadcast_node_to_edges.md): Broadcasts
+values from nodes to incident edges.
 
 [`check_compatible_with_schema_pb(...)`](./tfgnn/check_compatible_with_schema_pb.md):
 Checks that the given spec or value is compatible with the graph schema.
@@ -103,28 +117,47 @@ Checks that the given spec or value is compatible with the graph schema.
 [`check_homogeneous_graph_tensor(...)`](./tfgnn/check_homogeneous_graph_tensor.md):
 Raises ValueError when tfgnn.get_homogeneous_node_and_edge_set_name() does.
 
-[`check_required_features(...)`](./tfgnn/check_required_features.md): Checks the requirements of a given schema against another.
+[`check_required_features(...)`](./tfgnn/check_required_features.md): Checks the
+requirements of a given schema against another.
 
-[`check_scalar_graph_tensor(...)`](./tfgnn/check_scalar_graph_tensor.md)
+[`check_scalar_graph_tensor(...)`](./tfgnn/check_scalar_graph_tensor.md): Checks
+that graph tensor is scalar (has rank 0).
 
-[`combine_values(...)`](./tfgnn/combine_values.md): Combines a list of tensors into one (by concatenation or otherwise).
+[`combine_values(...)`](./tfgnn/combine_values.md): Combines a list of tensors
+into one (by concatenation or otherwise).
 
 [`convert_to_line_graph(...)`](./tfgnn/convert_to_line_graph.md): Obtain a
 graph's line graph.
 
-[`create_graph_spec_from_schema_pb(...)`](./tfgnn/create_graph_spec_from_schema_pb.md): Converts a graph schema proto message to a scalar GraphTensorSpec.
+[`create_graph_spec_from_schema_pb(...)`](./tfgnn/create_graph_spec_from_schema_pb.md):
+Converts a graph schema proto message to a scalar GraphTensorSpec.
 
 [`create_schema_pb_from_graph_spec(...)`](./tfgnn/create_schema_pb_from_graph_spec.md):
 Converts scalar GraphTensorSpec to a graph schema proto message.
 
-[`dataset_filter_with_summary(...)`](./tfgnn/dataset_filter_with_summary.md): Dataset filter with a summary for the fraction of dataset elements removed.
+[`dataset_filter_with_summary(...)`](./tfgnn/dataset_filter_with_summary.md):
+Dataset filter with a summary for the fraction of dataset elements removed.
 
 [`dataset_from_generator(...)`](./tfgnn/dataset_from_generator.md): Creates
 dataset from generator of any nest of scalar graph pieces.
 
-[`find_tight_size_constraints(...)`](./tfgnn/find_tight_size_constraints.md): Returns smallest possible size constraints that allow dataset padding.
+[`disable_graph_tensor_validation(...)`](./tfgnn/disable_graph_tensor_validation.md):
+Disables both static and runtime checks of graph tensors.
 
-[`gather_first_node(...)`](./tfgnn/gather_first_node.md): Gathers feature value from the first node of each graph component.
+[`disable_graph_tensor_validation_at_runtime(...)`](./tfgnn/disable_graph_tensor_validation_at_runtime.md):
+Disables runtime checks (`tf.debugging.Assert`) of graph tensors.
+
+[`enable_graph_tensor_validation(...)`](./tfgnn/enable_graph_tensor_validation.md):
+Enables static checks of graph tensors.
+
+[`enable_graph_tensor_validation_at_runtime(...)`](./tfgnn/enable_graph_tensor_validation_at_runtime.md):
+Enables both static and runtime checks of graph tensors.
+
+[`find_tight_size_constraints(...)`](./tfgnn/find_tight_size_constraints.md):
+Returns smallest possible size constraints that allow dataset padding.
+
+[`gather_first_node(...)`](./tfgnn/gather_first_node.md): Gathers feature value
+from the first node of each graph component.
 
 [`get_aux_type_prefix(...)`](./tfgnn/get_aux_type_prefix.md): Returns type
 prefix of aux node or edge set names, or `None` if non-aux.
@@ -132,11 +165,14 @@ prefix of aux node or edge set names, or `None` if non-aux.
 [`get_homogeneous_node_and_edge_set_name(...)`](./tfgnn/get_homogeneous_node_and_edge_set_name.md):
 Returns the sole `node_set_name, edge_set_name` or raises `ValueError`.
 
-[`get_io_spec(...)`](./tfgnn/get_io_spec.md): Returns tf.io parsing features for `GraphTensorSpec` type spec.
+[`get_io_spec(...)`](./tfgnn/get_io_spec.md): Returns tf.io parsing features for
+`GraphTensorSpec` type spec.
 
-[`get_registered_reduce_operation_names(...)`](./tfgnn/get_registered_reduce_operation_names.md): Returns the registered list of supported reduce operation names.
+[`get_registered_reduce_operation_names(...)`](./tfgnn/get_registered_reduce_operation_names.md):
+Returns the registered list of supported reduce operation names.
 
-[`graph_tensor_to_values(...)`](./tfgnn/graph_tensor_to_values.md): Convert an eager `GraphTensor` to a mapping of mappings of PODTs.
+[`graph_tensor_to_values(...)`](./tfgnn/graph_tensor_to_values.md): Convert an
+eager `GraphTensor` to a mapping of mappings of PODTs.
 
 [`homogeneous(...)`](./tfgnn/homogeneous.md): Constructs a homogeneous
 `GraphTensor` with node features and one edge_set.
@@ -144,16 +180,20 @@ Returns the sole `node_set_name, edge_set_name` or raises `ValueError`.
 [`is_dense_tensor(...)`](./tfgnn/is_dense_tensor.md): Returns whether a tensor
 (TF or Keras) is a Tensor.
 
-[`is_graph_tensor(...)`](./tfgnn/is_graph_tensor.md): Returns whether `value` is a GraphTensor (possibly wrapped for Keras).
+[`is_graph_tensor(...)`](./tfgnn/is_graph_tensor.md): Returns whether `value` is
+a GraphTensor (possibly wrapped for Keras).
 
 [`is_ragged_tensor(...)`](./tfgnn/is_ragged_tensor.md): Returns whether a tensor
 (TF or Keras) is a RaggedTensor.
 
-[`iter_features(...)`](./tfgnn/iter_features.md): Utility function to iterate over the features of a graph schema.
+[`iter_features(...)`](./tfgnn/iter_features.md): Utility function to iterate
+over the features of a graph schema.
 
-[`iter_sets(...)`](./tfgnn/iter_sets.md): Utility function to iterate over all the sets present in a graph schema.
+[`iter_sets(...)`](./tfgnn/iter_sets.md): Utility function to iterate over all
+the sets present in a graph schema.
 
-[`learn_fit_or_skip_size_constraints(...)`](./tfgnn/learn_fit_or_skip_size_constraints.md): Learns the optimal size constraints for the fixed size batching with retry.
+[`learn_fit_or_skip_size_constraints(...)`](./tfgnn/learn_fit_or_skip_size_constraints.md):
+Learns the optimal size constraints for the fixed size batching with retry.
 
 [`mask_edges(...)`](./tfgnn/mask_edges.md): Creates a GraphTensor after applying
 edge_mask over the specified edge-set.
@@ -161,35 +201,53 @@ edge_mask over the specified edge-set.
 [`node_degree(...)`](./tfgnn/node_degree.md): Returns the degree of each node
 w.r.t. one side of an edge set.
 
-[`pad_to_total_sizes(...)`](./tfgnn/pad_to_total_sizes.md): Pads graph tensor to the total sizes by inserting fake graph components.
+[`pad_to_total_sizes(...)`](./tfgnn/pad_to_total_sizes.md): Pads graph tensor to
+the total sizes by inserting fake graph components.
 
-[`parse_example(...)`](./tfgnn/parse_example.md): Parses a batch of serialized Example protos into a single `GraphTensor`.
+[`parse_example(...)`](./tfgnn/parse_example.md): Parses a batch of serialized
+Example protos into a single `GraphTensor`.
 
-[`parse_schema(...)`](./tfgnn/parse_schema.md): Parse a schema from text-formatted protos.
+[`parse_schema(...)`](./tfgnn/parse_schema.md): Parse a schema from
+text-formatted protos.
 
-[`parse_single_example(...)`](./tfgnn/parse_single_example.md): Parses a single serialized Example proto into a single `GraphTensor`.
+[`parse_single_example(...)`](./tfgnn/parse_single_example.md): Parses a single
+serialized Example proto into a single `GraphTensor`.
 
 [`pool(...)`](./tfgnn/pool.md): Pools values from edges to nodes, or from nodes
 or edges to context.
 
-[`pool_edges_to_context(...)`](./tfgnn/pool_edges_to_context.md): Aggregates (pools) edge values to graph context.
+[`pool_edges_to_context(...)`](./tfgnn/pool_edges_to_context.md): Aggregates
+(pools) edge values to graph context.
+
+[`pool_edges_to_node(...)`](./tfgnn/pool_edges_to_node.md): Aggregates (pools)
+edge values to incident nodes.
 
-[`pool_edges_to_node(...)`](./tfgnn/pool_edges_to_node.md): Aggregates (pools) edge values to incident nodes.
+[`pool_neighbors_to_node(...)`](./tfgnn/pool_neighbors_to_node.md): Aggregates
+(pools) neighbor node values along one or more edge sets.
 
-[`pool_nodes_to_context(...)`](./tfgnn/pool_nodes_to_context.md): Aggregates (pools) node values to graph context.
+[`pool_neighbors_to_node_feature(...)`](./tfgnn/pool_neighbors_to_node_feature.md):
+Aggregates (pools) sender node feature to receiver nodes feature.
 
-[`random_graph_tensor(...)`](./tfgnn/random_graph_tensor.md): Generate a graph tensor from a schema, with random features.
+[`pool_nodes_to_context(...)`](./tfgnn/pool_nodes_to_context.md): Aggregates
+(pools) node values to graph context.
 
-[`read_schema(...)`](./tfgnn/read_schema.md): Read a proto schema from a file with text-formatted contents.
+[`random_graph_tensor(...)`](./tfgnn/random_graph_tensor.md): Generate a graph
+tensor from a spec, with random features.
+
+[`read_schema(...)`](./tfgnn/read_schema.md): Read a proto schema from a file
+with text-formatted contents.
 
 [`reorder_nodes(...)`](./tfgnn/reorder_nodes.md): Reorders nodes within node
 sets according to indices.
 
-[`reverse_tag(...)`](./tfgnn/reverse_tag.md): Flips tfgnn.SOURCE to tfgnn.TARGET and vice versa.
+[`reverse_tag(...)`](./tfgnn/reverse_tag.md): Flips tfgnn.SOURCE to tfgnn.TARGET
+and vice versa.
 
-[`satisfies_size_constraints(...)`](./tfgnn/satisfies_size_constraints.md): Returns whether the input `graph_tensor` satisfies `total_sizes`.
+[`satisfies_size_constraints(...)`](./tfgnn/satisfies_size_constraints.md):
+Returns whether the input `graph_tensor` satisfies `total_sizes`.
 
-[`satisfies_total_sizes(...)`](./tfgnn/satisfies_size_constraints.md): Returns whether the input `graph_tensor` satisfies `total_sizes`.
+[`satisfies_total_sizes(...)`](./tfgnn/satisfies_size_constraints.md): Returns
+whether the input `graph_tensor` satisfies `total_sizes`.
 
 [`shuffle_features_globally(...)`](./tfgnn/shuffle_features_globally.md):
 Shuffles context, node set and edge set features of a scalar GraphTensor.
@@ -197,9 +255,11 @@ Shuffles context, node set and edge set features of a scalar GraphTensor.
 [`shuffle_nodes(...)`](./tfgnn/shuffle_nodes.md): Randomly reorders nodes of
 given node sets, within each graph component.
 
-[`softmax(...)`](./tfgnn/softmax.md): Computes softmax over a many-to-one relationship in a GraphTensor.
+[`softmax(...)`](./tfgnn/softmax.md): Computes softmax over a many-to-one
+relationship in a GraphTensor.
 
-[`softmax_edges_per_node(...)`](./tfgnn/softmax_edges_per_node.md): Returns softmax() of edge values per common `node_tag` node.
+[`softmax_edges_per_node(...)`](./tfgnn/softmax_edges_per_node.md): Returns
+softmax() of edge values per common `node_tag` node.
 
 [`structured_readout(...)`](./tfgnn/structured_readout.md): Reads out a feature
 value from select nodes (or edges) in a graph.
@@ -213,11 +273,14 @@ Checks `graph` supports `structured_readout()` from `required_keys`.
 [`validate_graph_tensor_spec_for_readout(...)`](./tfgnn/validate_graph_tensor_spec_for_readout.md):
 Checks `graph_spec` supports `structured_readout()` from `required_keys`.
 
-[`validate_schema(...)`](./tfgnn/validate_schema.md): Validates the correctness of a graph schema instance.
+[`validate_schema(...)`](./tfgnn/validate_schema.md): Validates the correctness
+of a graph schema instance.
 
-[`write_example(...)`](./tfgnn/write_example.md): Encode an eager `GraphTensor` to a tf.train.Example proto.
+[`write_example(...)`](./tfgnn/write_example.md): Encode an eager `GraphTensor`
+to a tf.train.Example proto.
 
-[`write_schema(...)`](./tfgnn/write_schema.md): Write a `GraphSchema` to a text-formatted proto file.
+[`write_schema(...)`](./tfgnn/write_schema.md): Write a `GraphSchema` to a
+text-formatted proto file.
 
 ## Type Aliases
 
@@ -233,9 +296,8 @@ Checks `graph_spec` supports `structured_readout()` from `required_keys`.
 
 [`IncidentNodeOrContextTag`](./tfgnn/IncidentNodeOrContextTag.md)
 
-
-
 
+
  
@@ -245,70 +307,70 @@ Checks `graph_spec` supports `structured_readout()` from `required_keys`.
 CONTEXT

 
 
Other Members

 
 

-`'context'`
+'context'
 
 

 
 EDGES
 
 
-`'edges'`
+'edges'
 
 

 
 HIDDEN_STATE
 
 
-`'hidden_state'`
+'hidden_state'
 
 

 
 NODES
 
 
-`'nodes'`
+'nodes'
 
 

 
 SIZE_NAME
 
 
-`'#size'`
+'#size'
 
 

 
 SOURCE
 
 
-`0`
+0
 
 

 
 SOURCE_NAME
 
 
-`'#source'`
+'#source'
 
 

 
 TARGET
 
 
-`1`
+1
 
 

 
 TARGET_NAME
 
 
-`'#target'`
+'#target'
 
 

 
 **version**
 
 
-`'0.6.0.dev1'`
+'1.0.0.dev2'
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/Adjacency.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/Adjacency.md
index f6caa09c..69c0c2c1 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/Adjacency.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/Adjacency.md
@@ -1,17 +1,10 @@
 # tfgnn.Adjacency
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Stores how edges connect pairs of nodes from source and target node sets.
 
@@ -19,103 +12,99 @@ Inherits From: [`HyperAdjacency`](../tfgnn/HyperAdjacency.md)
 
 
 tfgnn.Adjacency(
-    data: Data, spec: 'GraphPieceSpecBase', validate: bool = False
+    data: Data, spec: 'GraphPieceSpecBase'
 )
 

 
-
-
 
 
-Each hyper-edge connect one node from the source node set with one node from
-the target node sets. The source and target node sets could be the same.
-The adjacency information is a pair of integer tensors containing indices of
-nodes in source and target node sets. Those tensors are indexed by
-edges, have the same type spec and shape of `[*graph_shape, num_edges]`,
-where `num_edges` is the number of edges in the edge set (could be potentially
-ragged). The index tensors are of `tf.Tensor` type if `num_edges` is not
-`None` or `graph_shape.rank = 0` and of`tf.RaggedTensor` type otherwise.
+Each hyper-edge connect one node from the source node set with one node from the
+target node sets. The source and target node sets could be the same. The
+adjacency information is a pair of integer tensors containing indices of nodes
+in source and target node sets. Those tensors are indexed by edges, have the
+same type spec and shape of `[*graph_shape, num_edges]`, where `num_edges` is
+the number of edges in the edge set (could be potentially ragged). The index
+tensors are of `tf.Tensor` type if `num_edges` is not `None` or
+`graph_shape.rank = 0` and of`tf.RaggedTensor` type otherwise.
 
 The Adjacency is a composite tensor and a special case of tfgnn.HyperAdjacency
-class with tfgnn.SOURCE and tfgnn.TARGET node tags used for the source and
-target nodes correspondingly.
+class with tfgnn.SOURCE and
+tfgnn.TARGET node tags used for
+the source and target nodes correspondingly.
 
 
+
  
-
-
-

 
 
Args

 
 

 
-`data`
+data
 
 
 Nest of Field or subclasses of GraphPieceBase.
 
 

 
-`spec`
+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.
+A subclass of GraphPieceSpecBase with a _data_spec that matches
+data.
 
 

 

 
 
+
  
 
-
-

 
 
Attributes
  `indices_dtype`   The integer type
-to represent ragged splits.  
  `rank`  The rank of this Tensor. Guaranteed not to be `None`.  
 
-`shape`   A possibly-partial shape specification for
-this Tensor.
+
  indices_dtype   The
+dtype for graph items indexing. One of tf.int32 or
+tf.int64.  
  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 `TensorShape` is guaranteed to have a known rank, but the
+The returned TensorShape is guaranteed to have a known rank, but the
 individual dimension sizes may be unknown.
 
 

 
-`source`
+source
 
 
 The indices of source nodes.
 
 

 
-`source_name`
+source_name
 
 
 The node set name of source nodes.
 
 

 
-`spec`
+spec
 
 
 The public type specification of this tensor.
 
 

 
-`target`
+target
 
 
 The indices of target nodes.
 
 

 
-`target_name`
+target_name
 
 
 The node set name of target nodes.
@@ -127,7 +116,7 @@ The node set name of target nodes.
 
 
from_indices

 
-View
+View
 source
 
 
@@ -136,17 +125,14 @@ source
     source: tfgnn.Field,
     target: tfgnn.Field,
     *_,
-    validate: bool = True
+    validate: Optional[bool] = None
 ) -> 'Adjacency'
 

 
 Constructs a new instance from the `source` and `target` node indices.
 
-
 #### Example 1:
 
-
-
 ```python
 # Single graph (rank is 0). Connects pairs of nodes (a[0], b[2]),
 # (a[1], b[1]), (a[2], b[0]) from node sets a and b.
@@ -156,8 +142,6 @@ tfgnn.Adjacency.from_indices(('a', [0, 1, 2]),
 
 #### Example 2:
 
-
-
 ```python
 # Batch of two graphs (rank is 1). Connects pairs of nodes in
 # graph 0: (a[0], b[2]), (a[1], b[1]); graph 1: (a[2], b[0]).
@@ -166,60 +150,58 @@ tfgnn.Adjacency.from_indices(('a', tf.ragged.constant([[0, 1], [2]])),
 ```
 
 
+
  

 
 
Args

 
 

 
-`source`
+source
 
 
 The tuple of node set name and nodes index integer tensor. The
-index must have shape of `[*graph_shape, num_edges]`, where `num_edges`
+index must have shape of [*graph_shape, num_edges], where num_edges
 is the number of edges in each graph (could be ragged). It has
-`tf.Tensor` type if `num_edges` is not `None` or `graph_shape.rank = 0`
-and `tf.RaggedTensor` type otherwise.
+tf.Tensor type if num_edges is not None or graph_shape.rank = 0
+and tf.RaggedTensor type otherwise.
 
 

 
-`target`
+target
 
 
-Like `source` field, but for target edge endpoint. Index tensor
-must have the same type spec as for the `source`.
+Like source field, but for target edge endpoint. Index tensor
+must have the same type spec as for the source.
 
 

 
-`validate`
+validate
 
 
-If `True`, checks that source and target indices have the same
+If True, checks that source and target indices have the same
 type spec.
 
 

 

 
-
-
 
+
  

 
 
Returns

 

 
-An `Adjacency` tensor with a shape and an indices_dtype being inferred
-from the `indices` values.
+An Adjacency tensor with a shape and an indices_dtype being inferred
+from the indices values.
 
 

 
 

 
-
-
 
get_indices_dict

 
-View
+View
 source
 
 
@@ -228,10 +210,9 @@ source
 
 Returns copy of indices as a dictionary.
 
-
 
node_set_name

 
-View
+View
 source
 
 
@@ -242,24 +223,61 @@ source
 
 Returns a node set name for the given node set tag.
 
-
 
set_shape

 
-View
+View
 source
 
 
 set_shape(
     new_shape: ShapeLike
-) -> 'GraphPieceSpecBase'
+) -> 'GraphPieceBase'
 

 
-Enforce the common prefix shape on all the contained features.
+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
+View
 source
 
 
@@ -269,7 +287,3 @@ source
 

 
 Returns an index tensor for the given node set tag.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/AdjacencySpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/AdjacencySpec.md
index c1a33161..6c84b62f 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/AdjacencySpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/AdjacencySpec.md
@@ -1,19 +1,13 @@
 # tfgnn.AdjacencySpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
-A type spec for tfgnn.Adjacency.
+A type spec for
+tfgnn.Adjacency.
 
 Inherits From: [`HyperAdjacencySpec`](../tfgnn/HyperAdjacencySpec.md)
 
@@ -22,42 +16,43 @@ Inherits From: [`HyperAdjacencySpec`](../tfgnn/HyperAdjacencySpec.md)
     data_spec: DataSpec,
     shape: tf.TensorShape,
     indices_dtype: tf.dtypes.DType,
+    row_splits_dtype: tf.dtypes.DType,
     metadata: Metadata = None,
-    validate: bool = False
+    check_consistent_indices_dtype: bool = False,
+    check_consistent_row_splits_dtype: bool = False
 )
 

 
-
-
 
-
-
-
-
 
+
  
 
-
-
-

 
 
Attributes
  `indices_dtype`   The integer type
-to represent ragged splits.  
  `rank`  The rank of the GraphPiece. Guaranteed not to be `None`.  
 `shape`   A possibly-partial shape specification
-of the GraphPiece.
-
-The returned `TensorShape` is guaranteed to have a known rank, but the
-individual dimension sizes may be unknown.  
 
-`source`  
-
- 
  `source_name`   Returns
-the node set name for source nodes.  
 
-`target`  
-
- 
  `target_name`   Returns
-the node set name for target nodes.  
 
-`total_size`   The total number of edges if
-known.  
  `value_type`   The
-Python type for values that are compatible with this TypeSpec.
+
  indices_dtype   The
+dtype for graph items indexing. One of tf.int32 or
+tf.int64.  
  rank
+  The rank of the GraphPiece. 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 of the GraphPiece.
+
+The returned TensorShape is guaranteed to have a known rank, but
+the individual dimension sizes may be unknown.  
 
+source  
+
+ 
  source_name  
+Returns the node set name for source nodes.  
 
+target  
+
+ 
  target_name  
+Returns the node set name for target nodes.  
 
+total_size   The total number of
+edges if known.  
 
+value_type   The Python type for
+values that are compatible with this TypeSpec.
 
 In particular, all values that are compatible with this TypeSpec must be an
 instance of this type.
@@ -65,8 +60,6 @@ instance of this type.
 

 

 
-
-
 ## Methods
 
 
experimental_as_proto

@@ -93,13 +86,14 @@ Returns a TypeSpec instance based on the serialized proto.
 Do NOT override for custom non-TF types.
 
 
+
  

 
 
Args

 
 

 
-`proto`
+proto
 
 
 Proto generated using 'experimental_as_proto'.
@@ -120,7 +114,7 @@ Do NOT override for custom non-TF types.
 
 
from_incident_node_sets

 
-View
+View
 source
 
 
@@ -134,59 +128,56 @@ source
 
 Constructs a new instance from the `incident_node_sets`.
 
-
 
+
  

 
 
Args

 
 

 
-`source_node_set`
+source_node_set
 
 
 The name of the source node set.
 
 

 
-`target_node_set`
+target_node_set
 
 
 The name of the target node set.
 
 

 
-`index_spec`
+index_spec
 
 
 type spec for source and target index tensors of shape
-`[*graph_shape, num_edges]`, where num_edges is the number of edges in
-each graph. If `num_edges` is not `None` or `graph_shape.rank = 0` the
-spec must be of `tf.TensorSpec` type and of `tf.RaggedTensorSpec` type
+[*graph_shape, num_edges], where num_edges is the number of edges in
+each graph. If num_edges is not None or graph_shape.rank = 0 the
+spec must be of tf.TensorSpec type and of tf.RaggedTensorSpec type
 otherwise.
 
 

 

 
-
-
 
+
  

 
 
Returns

 

 
-A `AdjacencySpec` TypeSpec.
+A AdjacencySpec TypeSpec.
 
 

 
 

 
-
-
 
from_value

 
-View
+View
 source
 
 
@@ -198,10 +189,9 @@ source
 
 Extension Types API: Factory method.
 
-
 
get_index_specs_dict

 
-View
+View
 source
 
 
@@ -210,7 +200,6 @@ source
 
 Returns copy of indices type specs as a dictionary.
 
-
 
is_compatible_with

 
 
@@ -225,13 +214,14 @@ Prefer using "is_subtype_of" and "most_specific_common_supertype" wherever
 possible.
 
 
+
  

 
 
Args

 
 

 
-`spec_or_value`
+spec_or_value
 
 
 A TypeSpec or TypeSpec associated value to compare against.
@@ -239,8 +229,6 @@ A TypeSpec or TypeSpec associated value to compare against.
 

 

 
-
-
 
is_subtype_of

 
 
@@ -253,19 +241,19 @@ Returns True if `self` is a subtype of `other`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
 A TraceType object.
@@ -273,8 +261,6 @@ A TraceType object.
 

 

 
-
-
 
most_specific_common_supertype

 
 
@@ -283,23 +269,23 @@ A TraceType object.
 ) -> Optional['TypeSpec']
 

 
-Returns the most specific supertype TypeSpec  of `self` and `others`.
+Returns the most specific supertype TypeSpec of `self` and `others`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`others`
+others
 
 
 A sequence of TraceTypes.
@@ -307,8 +293,6 @@ A sequence of TraceTypes.
 

 

 
-
-
 
most_specific_compatible_type

 
 
@@ -317,53 +301,51 @@ A sequence of TraceTypes.
 ) -> 'TypeSpec'
 

 
-Returns the most specific TypeSpec compatible with `self` and `other`. (deprecated)
+Returns the most specific TypeSpec compatible with `self` and `other`.
+(deprecated)
 
 Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version.
-Instructions for updating:
-Use most_specific_common_supertype instead.
+Instructions for updating: Use most_specific_common_supertype instead.
 
-Deprecated. Please use `most_specific_common_supertype` instead.
-Do not override this function.
+Deprecated. Please use `most_specific_common_supertype` instead. Do not override
+this function.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
-A `TypeSpec`.
+A TypeSpec.
 
 

 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If there is no TypeSpec that is compatible with both `self`
-and `other`.
+If there is no TypeSpec that is compatible with both self
+and other.
 
 

 

 
-
-
 
node_set_name

 
-View
+View
 source
 
 
@@ -376,7 +358,7 @@ Returns a node set name for the given node set tag.
 
 
relax

 
-View
+View
 source
 
 
@@ -390,13 +372,14 @@ Allows variable number of graph edges.
 Calling with all default parameters keeps the spec unchanged.
 
 
+
  

 
 
Args

 
 

 
-`num_edges`
+num_edges
 
 
 if True, allows a variable number of edges in each edge set.
@@ -406,6 +389,7 @@ If False, returns spec unchanged.
 

 
 
+
  
@@ -418,13 +402,14 @@ Relaxed compatible spec.
 

 
 
Returns

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if adjacency is not scalar (rank > 0).
@@ -432,6 +417,45 @@ if adjacency is not scalar (rank > 0).
 

 

 
+
with_indices_dtype

+
+View
+source
+
+
+with_indices_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_row_splits_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_shape(
+    new_shape: ShapeLike
+) -> 'GraphPieceSpecBase'
+

+
+Enforce the common prefix shape on all the contained features.
+
 
__eq__

 
 
@@ -442,10 +466,9 @@ if adjacency is not scalar (rank > 0).
 
 Return self==value.
 
-
 
__getitem__

 
-View
+View
 source
 
 
@@ -456,7 +479,6 @@ source
 
 Returns an index tensor type spec for the given node set tag.
 
-
 
__ne__

 
 
@@ -466,7 +488,3 @@ Returns an index tensor type spec for the given node set tag.
 

 
 Return self!=value.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md
index 7172739c..c978edf8 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md
@@ -1,110 +1,98 @@
 # tfgnn.Context
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A composite tensor for graph context features.
 
 
 tfgnn.Context(
-    data: Data, spec: 'GraphPieceSpecBase', validate: bool = False
+    data: Data, spec: 'GraphPieceSpecBase'
 )
 

 
-
-
 
 
-The items of the context are the graph components (just like the items of a
-node set are the nodes and the items of an edge set are the edges). The
-`Context` is a composite tensor. It stores features that belong to a graph
-component as a whole, not any particular node or edge. Each context feature
-has a shape `[*graph_shape, num_components, ...]`, where `num_components` is
-the number of graph components in a graph (could be ragged).
+The items of the context are the graph components (just like the items of a node
+set are the nodes and the items of an edge set are the edges). The `Context` is
+a composite tensor. It stores features that belong to a graph component as a
+whole, not any particular node or edge. Each context feature has a shape
+`[*graph_shape, num_components, ...]`, where `num_components` is the number of
+graph components in a graph (could be ragged).
 
 
+
  
-
-
-

 
 
Args

 
 

 
-`data`
+data
 
 
 Nest of Field or subclasses of GraphPieceBase.
 
 

 
-`spec`
+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.
+A subclass of GraphPieceSpecBase with a _data_spec that matches
+data.
 
 

 

 
 
+
  
 
-
-

 
 
Attributes
  `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
+
  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 TensorShape is guaranteed to have a known rank, but the
 individual dimension sizes may be unknown.
 
 

 
-`sizes`
+sizes
 
 
 The number of items in each graph component.
 
 

 
-`spec`
+spec
 
 
 The public type specification of this tensor.
 
 

 
-`total_num_components`
+total_num_components
 
 
 The total number of graph components.
 
 

 
-`total_size`
+total_size
 
 
 The total number of items.
@@ -116,97 +104,105 @@ The total number of items.
 
 
from_fields

 
-View
+View
 source
 
 
 @classmethod
 from_fields(
-    *,
+    *_,
     features: Optional[Fields] = None,
     sizes: Optional[Field] = None,
     shape: Optional[ShapeLike] = None,
-    indices_dtype: Optional[tf.dtypes.DType] = None
+    indices_dtype: Optional[tf.dtypes.DType] = None,
+    validate: Optional[bool] = None
 ) -> 'Context'
 

 
 Constructs a new instance from context fields.
 
-
 #### Example:
 
-
-
 ```python
 tfgnn.Context.from_fields(features={'country_code': ['CH']})
 ```
 
 
+
  
+
+
+

 
 
Args

 
 

 
-`features`
+features
 
 
 A mapping from feature name to feature Tensor or RaggedTensor.
-All feature tensors must have shape `[*graph_shape, num_components,
-*feature_shape]`, where `num_components` is the number of graph
-components (could be ragged); `feature_shape` are feature-specific inner
+All feature tensors must have shape [*graph_shape, num_components,
+*feature_shape], where num_components is the number of graph
+components (could be ragged); feature_shape are feature-specific inner
 dimensions.
 
 

 
-`sizes`
+sizes
 
 
-A Tensor of 1's with shape `[*graph_shape, num_components]`, where
-`num_components` is the number of graph components (could be ragged).
-For symmetry with `sizes` in NodeSet and EdgeSet, this counts the items
+A Tensor of 1's with shape [*graph_shape, num_components], where
+num_components is the number of graph components (could be ragged).
+For symmetry with sizes in NodeSet and EdgeSet, this counts the items
 per graph component, but since the items of Context are the components
-themselves, each value is 1. Must be compatible with `shape`, if that is
+themselves, each value is 1. Must be compatible with shape, if that is
 specified.
 
 

 
-`shape`
+shape
 
 
 The shape of this tensor and a GraphTensor containing it, also
-known as the `graph_shape`. If not specified, the shape is inferred from
-`sizes` or set to `[]` if the `sizes` is not specified.
+known as the graph_shape. If not specified, the shape is inferred from
+sizes or set to [] if the sizes is not specified.
 
 

 
-`indices_dtype`
+indices_dtype
 
 
-An `indices_dtype` of a GraphTensor containing this object,
-used as `row_splits_dtype` when batching potentially ragged fields. If
-`sizes` are specified they are casted to that type.
+An indices_dtype of a GraphTensor containing this object,
+used as row_splits_dtype when batching potentially ragged fields. If
+sizes are specified they are casted to that type.
+

+validate
+
+If true, use tf.assert ops to inspect the shapes of each field
+and check at runtime that they form a valid Context. The default
+behavior is set by the disable_graph_tensor_validation_at_runtime()
+and enable_graph_tensor_validation_at_runtime().
 
 

 

 
 
+
  

 
 
Returns

 

 
-A `Context` composite tensor.
+A Context composite tensor.
 
 

 
 

 
-
-
 
get_features_dict

 
-View
+View
 source
 
 
@@ -215,10 +211,9 @@ source
 
 Returns features copy as a dictionary.
 
-
 
replace_features

 
-View
+View
 source
 
 
@@ -229,20 +224,57 @@ source
 
 Returns a new instance with a new set of features.
 
-
 
set_shape

 
-View
+View
 source
 
 
 set_shape(
     new_shape: ShapeLike
-) -> 'GraphPieceSpecBase'
+) -> 'GraphPieceBase'
 

 
-Enforce the common prefix shape on all the contained features.
+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__

 
@@ -256,7 +288,3 @@ source
 

 
 Indexing operator `[]` to access feature values by their name.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/ContextSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/ContextSpec.md
index 035b4e88..aa470448 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/ContextSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/ContextSpec.md
@@ -1,17 +1,10 @@
 # tfgnn.ContextSpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A type spec for tfgnn.Context.
 
@@ -20,40 +13,41 @@ A type spec for tfgnn.Context.
     data_spec: DataSpec,
     shape: tf.TensorShape,
     indices_dtype: tf.dtypes.DType,
+    row_splits_dtype: tf.dtypes.DType,
     metadata: Metadata = None,
-    validate: bool = False
+    check_consistent_indices_dtype: bool = False,
+    check_consistent_row_splits_dtype: bool = False
 )
 

 
-
-
 
-
-
-
-
 
+
  
 
-

 
 
Attributes
  `features_spec`   A read-only
-mapping of feature name to feature spec.  
 
-`indices_dtype`   The integer type to
-represent ragged splits.  
  `rank`  
-The rank of the GraphPiece. Guaranteed not to be `None`.  
 
-`shape`   A possibly-partial shape specification of
-the GraphPiece.
-
-The returned `TensorShape` is guaranteed to have a known rank, but the
-individual dimension sizes may be unknown.  
 
-`sizes_spec`   The type spec for the sizes that
-provides num. elements per component.  
 
-`total_num_components`   The total
-number of graph components if known.  
 
-`total_size`   The total number of graph items
-if known.  
  `value_type`  
-The Python type for values that are compatible with this TypeSpec.
+
  features_spec   A
+read-only mapping of feature name to feature spec.  
 
+indices_dtype   The dtype for
+graph items indexing. One of tf.int32 or tf.int64.
+ 
  rank   The rank of
+the GraphPiece. 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 of the GraphPiece.
+
+The returned TensorShape is guaranteed to have a known rank, but
+the individual dimension sizes may be unknown.  
 
+sizes_spec   The type spec for the
+sizes that provides num. elements per component.  
 
+total_num_components  
+The total number of graph components if known.  
 
+total_size   The total number of
+graph items if known.  
 
+value_type   The Python type for
+values that are compatible with this TypeSpec.
 
 In particular, all values that are compatible with this TypeSpec must be an
 instance of this type.
@@ -61,8 +55,6 @@ instance of this type.
 

 

 
-
-
 ## Methods
 
 
experimental_as_proto

@@ -89,13 +81,14 @@ Returns a TypeSpec instance based on the serialized proto.
 Do NOT override for custom non-TF types.
 
 
+
  

 
 
Args

 
 

 
-`proto`
+proto
 
 
 Proto generated using 'experimental_as_proto'.
@@ -116,7 +109,7 @@ Do NOT override for custom non-TF types.
 
 
from_field_specs

 
-View
+View
 source
 
 
@@ -130,12 +123,13 @@ source
 ) -> 'ContextSpec'
 

 
-The counterpart of Context.from_fields() for field type specs.
-
+The counterpart of
+Context.from_fields()
+for field type specs.
 
 
from_value

 
-View
+View
 source
 
 
@@ -147,7 +141,6 @@ source
 
 Extension Types API: Factory method.
 
-
 
is_compatible_with

 
 
@@ -162,13 +155,14 @@ Prefer using "is_subtype_of" and "most_specific_common_supertype" wherever
 possible.
 
 
+
  

 
 
Args

 
 

 
-`spec_or_value`
+spec_or_value
 
 
 A TypeSpec or TypeSpec associated value to compare against.
@@ -176,8 +170,6 @@ A TypeSpec or TypeSpec associated value to compare against.
 

 

 
-
-
 
is_subtype_of

 
 
@@ -190,19 +182,19 @@ Returns True if `self` is a subtype of `other`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
 A TraceType object.
@@ -210,8 +202,6 @@ A TraceType object.
 

 

 
-
-
 
most_specific_common_supertype

 
 
@@ -220,23 +210,23 @@ A TraceType object.
 ) -> Optional['TypeSpec']
 

 
-Returns the most specific supertype TypeSpec  of `self` and `others`.
+Returns the most specific supertype TypeSpec of `self` and `others`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`others`
+others
 
 
 A sequence of TraceTypes.
@@ -244,8 +234,6 @@ A sequence of TraceTypes.
 

 

 
-
-
 
most_specific_compatible_type

 
 
@@ -254,51 +242,51 @@ A sequence of TraceTypes.
 ) -> 'TypeSpec'
 

 
-Returns the most specific TypeSpec compatible with `self` and `other`. (deprecated)
+Returns the most specific TypeSpec compatible with `self` and `other`.
+(deprecated)
 
 Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version.
-Instructions for updating:
-Use most_specific_common_supertype instead.
+Instructions for updating: Use most_specific_common_supertype instead.
 
-Deprecated. Please use `most_specific_common_supertype` instead.
-Do not override this function.
+Deprecated. Please use `most_specific_common_supertype` instead. Do not override
+this function.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
-A `TypeSpec`.
+A TypeSpec.
 
 

 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If there is no TypeSpec that is compatible with both `self`
-and `other`.
+If there is no TypeSpec that is compatible with both self
+and other.
 
 

 

 
 
relax

 
-View
+View
 source
 
 
@@ -312,22 +300,24 @@ Allows variable number of graph components.
 Calling with all default parameters keeps the spec unchanged.
 
 
+
  

 
 
Args

 
 

 
-`num_components`
+num_components
 
 
 if True, allows variable number of graph components by
-setting the outermost sizes dimension to `None`.
+setting the outermost sizes dimension to None.
 
 

 

 
 
+
  
@@ -340,13 +330,14 @@ Relaxed compatible context spec.
 

 
 
Returns

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if contex is not scalar (rank > 0).
@@ -354,6 +345,45 @@ if contex is not scalar (rank > 0).
 

 

 
+
with_indices_dtype

+
+View
+source
+
+
+with_indices_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_row_splits_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_shape(
+    new_shape: ShapeLike
+) -> 'GraphPieceSpecBase'
+

+
+Enforce the common prefix shape on all the contained features.
+
 
__eq__

 
 
@@ -364,10 +394,9 @@ if contex is not scalar (rank > 0).
 
 Return self==value.
 
-
 
__getitem__

 
-View
+View
 source
 
 
@@ -385,7 +414,3 @@ source
 

 
 Return self!=value.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md
index 907bafce..a4856143 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md
@@ -1,36 +1,28 @@
 # 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
+    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).
+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,
+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
@@ -38,84 +30,82 @@ 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.
+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`
+data
 
 
 Nest of Field or subclasses of GraphPieceBase.
 
 

 
-`spec`
+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.
+A subclass of GraphPieceSpecBase with a _data_spec that matches
+data.
 
 

 

 
 
+
  
 
-
-
+
@@ -91,7 +84,7 @@ ops from this function into a dependency chain.
 
 

 
 
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
+
  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 TensorShape is guaranteed to have a known rank, but the
 individual dimension sizes may be unknown.
 
 

 
-`sizes`
+sizes
 
 
 The number of items in each graph component.
 
 

 
-`spec`
+spec
 
 
 The public type specification of this tensor.
 
 

 
-`total_num_components`
+total_num_components
 
 
 The total number of graph components.
 
 

 
-`total_size`
+total_size
 
 
 The total number of items.
@@ -127,19 +117,22 @@ The total number of items.
 
 
from_fields

 
-View
+View
 source
 
 
 @classmethod
 from_fields(
-    *, features: Optional[Fields] = None, sizes: Field, adjacency: Adjacency
+    *_,
+    features: Optional[Fields] = None,
+    sizes: Field,
+    adjacency: Adjacency,
+    validate: Optional[bool] = None
 ) -> 'EdgeSet'
 

 
 Constructs a new instance from edge set fields.
 
-
 #### Example 1:
 
 ```python
@@ -161,57 +154,67 @@ tfgnn.EdgeSet.from_fields(
 ```
 
 
+
  
+
+
+

 
 
Args

 
 

 
-`features`
+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
+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`
+sizes
 
 
 The number of edges in each graph component. Has shape
-`[*graph_shape, num_components]`, where `num_components` is the number
+[*graph_shape, num_components], where num_components is the number
 of graph components (could be ragged).
 
 

 
-`adjacency`
+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.
+An EdgeSet composite tensor.
 
 

 
 

 
-
-
 
get_features_dict

 
-View
+View
 source
 
 
@@ -220,10 +223,9 @@ source
 
 Returns features copy as a dictionary.
 
-
 
replace_features

 
-View
+View
 source
 
 
@@ -234,20 +236,57 @@ source
 
 Returns a new instance with a new set of features.
 
-
 
set_shape

 
-View
+View
 source
 
 
 set_shape(
     new_shape: ShapeLike
-) -> 'GraphPieceSpecBase'
+) -> 'GraphPieceBase'
 

 
-Enforce the common prefix shape on all the contained features.
+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__

 
@@ -261,7 +300,3 @@ source
 

 
 Indexing operator `[]` to access feature values by their name.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSetSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSetSpec.md
index 97ec4f55..6c323a9f 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSetSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSetSpec.md
@@ -1,17 +1,10 @@
 # tfgnn.EdgeSetSpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A type spec for tfgnn.EdgeSet.
 
@@ -20,42 +13,43 @@ A type spec for tfgnn.EdgeSet.
     data_spec: DataSpec,
     shape: tf.TensorShape,
     indices_dtype: tf.dtypes.DType,
+    row_splits_dtype: tf.dtypes.DType,
     metadata: Metadata = None,
-    validate: bool = False
+    check_consistent_indices_dtype: bool = False,
+    check_consistent_row_splits_dtype: bool = False
 )
 

 
-
-
 
-
-
-
-
 
+
  
 
-

 
 
Attributes
  `adjacency_spec`   A type spec for
-the adjacency composite tensor.  
 
-`features_spec`   A read-only mapping of
-feature name to feature spec.  
 
-`indices_dtype`   The integer type to
-represent ragged splits.  
  `rank`  
-The rank of the GraphPiece. Guaranteed not to be `None`.  
 
-`shape`   A possibly-partial shape specification of
-the GraphPiece.
-
-The returned `TensorShape` is guaranteed to have a known rank, but the
-individual dimension sizes may be unknown.  
 
-`sizes_spec`   The type spec for the sizes that
-provides num. elements per component.  
 
-`total_num_components`   The total
-number of graph components if known.  
 
-`total_size`   The total number of edges if
-known.  
  `value_type`   The
-Python type for values that are compatible with this TypeSpec.
+
  adjacency_spec   A
+type spec for the adjacency composite tensor.  
 
+features_spec   A read-only
+mapping of feature name to feature spec.  
 
+indices_dtype   The dtype for
+graph items indexing. One of tf.int32 or tf.int64.
+ 
  rank   The rank of
+the GraphPiece. 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 of the GraphPiece.
+
+The returned TensorShape is guaranteed to have a known rank, but
+the individual dimension sizes may be unknown.  
 
+sizes_spec   The type spec for the
+sizes that provides num. elements per component.  
 
+total_num_components  
+The total number of graph components if known.  
 
+total_size   The total number of
+edges if known.  
 
+value_type   The Python type for
+values that are compatible with this TypeSpec.
 
 In particular, all values that are compatible with this TypeSpec must be an
 instance of this type.
@@ -63,8 +57,6 @@ instance of this type.
 

 

 
-
-
 ## Methods
 
 
experimental_as_proto

@@ -91,13 +83,14 @@ Returns a TypeSpec instance based on the serialized proto.
 Do NOT override for custom non-TF types.
 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_for_readout.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_for_readout.md
index 64accbf6..53bb11fe 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_for_readout.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_for_readout.md
@@ -1,17 +1,10 @@
 # tfgnn.validate_graph_tensor_for_readout
 
-[TOC]
-
 
 
-

 
 
Args

 
 

 
-`proto`
+proto
 
 
 Proto generated using 'experimental_as_proto'.
@@ -118,7 +111,7 @@ Do NOT override for custom non-TF types.
 
 
from_field_specs

 
-View
+View
 source
 
 
@@ -131,12 +124,13 @@ source
 ) -> 'EdgeSetSpec'
 

 
-The counterpart of EdgeSet.from_fields() for values type specs.
-
+The counterpart of
+EdgeSet.from_fields()
+for values type specs.
 
 
from_value

 
-View
+View
 source
 
 
@@ -148,7 +142,6 @@ source
 
 Extension Types API: Factory method.
 
-
 
is_compatible_with

 
 
@@ -163,13 +156,14 @@ Prefer using "is_subtype_of" and "most_specific_common_supertype" wherever
 possible.
 
 
+
  

 
 
Args

 
 

 
-`spec_or_value`
+spec_or_value
 
 
 A TypeSpec or TypeSpec associated value to compare against.
@@ -177,8 +171,6 @@ A TypeSpec or TypeSpec associated value to compare against.
 

 

 
-
-
 
is_subtype_of

 
 
@@ -191,19 +183,19 @@ Returns True if `self` is a subtype of `other`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
 A TraceType object.
@@ -211,8 +203,6 @@ A TraceType object.
 

 

 
-
-
 
most_specific_common_supertype

 
 
@@ -221,23 +211,23 @@ A TraceType object.
 ) -> Optional['TypeSpec']
 

 
-Returns the most specific supertype TypeSpec  of `self` and `others`.
+Returns the most specific supertype TypeSpec of `self` and `others`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`others`
+others
 
 
 A sequence of TraceTypes.
@@ -245,8 +235,6 @@ A sequence of TraceTypes.
 

 

 
-
-
 
most_specific_compatible_type

 
 
@@ -255,51 +243,51 @@ A sequence of TraceTypes.
 ) -> 'TypeSpec'
 

 
-Returns the most specific TypeSpec compatible with `self` and `other`. (deprecated)
+Returns the most specific TypeSpec compatible with `self` and `other`.
+(deprecated)
 
 Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version.
-Instructions for updating:
-Use most_specific_common_supertype instead.
+Instructions for updating: Use most_specific_common_supertype instead.
 
-Deprecated. Please use `most_specific_common_supertype` instead.
-Do not override this function.
+Deprecated. Please use `most_specific_common_supertype` instead. Do not override
+this function.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
-A `TypeSpec`.
+A TypeSpec.
 
 

 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If there is no TypeSpec that is compatible with both `self`
-and `other`.
+If there is no TypeSpec that is compatible with both self
+and other.
 
 

 

 
 
relax

 
-View
+View
 source
 
 
@@ -313,30 +301,32 @@ Allows variable number of edge or/and graph components.
 Calling with all default parameters keeps the spec unchanged.
 
 
+
  

 
 
Args

 
 

 
-`num_components`
+num_components
 
 
 if True, allows variable number of graph components by
-setting the outermost sizes dimension to `None`.
+setting the outermost sizes dimension to None.
 
 

 
-`num_edges`
+num_edges
 
 
 if True, allows variable number of edges by setting the
-outermost features dimensions to `None`.
+outermost features dimensions to None.
 
 

 

 
 
+
  
@@ -349,13 +339,14 @@ Relaxed compatible edge set spec.
 

 
 
Returns

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if edge set is not scalar (rank > 0).
@@ -363,6 +354,45 @@ if edge set is not scalar (rank > 0).
 

 

 
+
with_indices_dtype

+
+View
+source
+
+
+with_indices_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_row_splits_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_shape(
+    new_shape: ShapeLike
+) -> 'GraphPieceSpecBase'
+

+
+Enforce the common prefix shape on all the contained features.
+
 
__eq__

 
 
@@ -373,10 +403,9 @@ if edge set is not scalar (rank > 0).
 
 Return self==value.
 
-
 
__getitem__

 
-View
+View
 source
 
 
@@ -394,7 +423,3 @@ source
 

 
 Return self!=value.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/Feature.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/Feature.md
deleted file mode 100644
index db5d2143..00000000
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/Feature.md
+++ /dev/null
@@ -1,75 +0,0 @@
-# tfgnn.Feature
-
-[TOC]
-
-
-
-
-
-
-
-
-
-A schema for a single feature.
-
-
-
-This proto message contains the description, shape, data type and some more
-fields about a feature in the schema.
-
-
-
-
- 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
Attributes

-`description`
-
-`string description`
-

-`dtype`
-
-`DataType dtype`
-

-`example_values`
-
-`repeated Feature example_values`
-

-`sample_values`
-
-`Feature sample_values`
-

-`shape`
-
-`TensorShapeProto shape`
-

-`source`
-
-`string source`
-

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/FeatureDefaultValues.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/FeatureDefaultValues.md
index 825fcc85..3f14d743 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/FeatureDefaultValues.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/FeatureDefaultValues.md
@@ -1,17 +1,10 @@
 # tfgnn.FeatureDefaultValues
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Default values for graph context, node sets and edge sets features.
 
@@ -21,38 +14,33 @@ Default values for graph context, node sets and edge sets features.
 )
 

 
-
-
 
-
-
-
-
 
+
  

 
 
Attributes

 
 

 
-`context`
+context
 
 
-A `namedtuple` alias for field number 0
+A namedtuple alias for field number 0
 
 

 
-`node_sets`
+node_sets
 
 
-A `namedtuple` alias for field number 1
+A namedtuple alias for field number 1
 
 

 
-`edge_sets`
+edge_sets
 
 
-A `namedtuple` alias for field number 2
+A namedtuple alias for field number 2
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/Field.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/Field.md
index 3a0e8272..a3ca2ac7 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/Field.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/Field.md
@@ -1,14 +1,8 @@
-

-
-
-

-
 # tfgnn.Field
 
 
-This symbol is a **type alias**.
-
 
+This symbol is a **type alias**.
 
 #### Source:
 
@@ -19,6 +13,4 @@ This symbol is a **type alias**.
 ]
 

 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldOrFields.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldOrFields.md
index e97514b7..f36b6f8a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldOrFields.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldOrFields.md
@@ -1,14 +1,8 @@
-

-
-
-

-
 # tfgnn.FieldOrFields
 
 
-This symbol is a **type alias**.
-
 
+This symbol is a **type alias**.
 
 #### Source:
 
@@ -20,6 +14,4 @@ This symbol is a **type alias**.
 ]
 

 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldSpec.md
index 45b5e93b..ec8ef6e0 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldSpec.md
@@ -1,14 +1,8 @@
-

-
-
-

-
 # tfgnn.FieldSpec
 
 
-This symbol is a **type alias**.
-
 
+This symbol is a **type alias**.
 
 #### Source:
 
@@ -19,6 +13,4 @@ This symbol is a **type alias**.
 ]
 

 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/Fields.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/Fields.md
index 76fa1f45..103402ae 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/Fields.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/Fields.md
@@ -1,14 +1,8 @@
-

-
-
-

-
 # tfgnn.Fields
 
 
-This symbol is a **type alias**.
-
 
+This symbol is a **type alias**.
 
 #### Source:
 
@@ -19,6 +13,4 @@ This symbol is a **type alias**.
 ]
 

 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldsSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldsSpec.md
index 628c9b06..cf61524a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldsSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/FieldsSpec.md
@@ -1,14 +1,8 @@
-

-
-
-

-
 # tfgnn.FieldsSpec
 
 
-This symbol is a **type alias**.
-
 
+This symbol is a **type alias**.
 
 #### Source:
 
@@ -19,6 +13,4 @@ This symbol is a **type alias**.
 ]
 
 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphSchema.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphSchema.md
deleted file mode 100644
index 472dd033..00000000
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphSchema.md
+++ /dev/null
@@ -1,90 +0,0 @@
-# tfgnn.GraphSchema
-
-[TOC]
-
-
-
-
-
-
-
-
-
-A schema definition for graphs.
-
-
-
-The `GraphSchema` message describes the sets of nodes and edges provided by the
-graph data structure. It also provides the lists of available features within
-each node set, each edge set and over each graph ("context" features).
-
-The purpose of this schema is to describe the tensors available in a graph data
-structure container and the associations between them, their data types, shapes
-and feature descriptions. It exists to allow the user to declare the topology of
-the graph, and for the I/O routines to automatically parse the data. It also
-contains metadata about the graph's various features. See the "Describing your
-Graph" section of the documentation for full details on how to create an
-instance of this message to define the shape and encoding of your dataset.
-
-Note that the schema does not provide a full definition for how the features are
-intended to be used or combined; that belongs to a separate description for a
-model, which is beyond the scope of this schema object. For instance, whether a
-context feature is broadcast over a particular node feature during learning is
-information that isn't related to the data stored in the container of features.
-
-#### Intended usage:
-
-
-
-* To accompany a graph container data structure, as documentation reporting
-  entities, edges and features available during training.
-* To be serialized in the metadata of training data files.
-* To be safeguarded along with model checkpoints in order to keep track of input
-  features used historically.
-* To be utilized to automatically infer good default models.
-
-Note that a feature names beginnning with `#` are explicitly reserved and
-disallowed. (These are used in serialization.)
-
-
-
-
- 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
Attributes

-`context`
-
-`Context context`
-

-`edge_sets`
-
-`repeated EdgeSetsEntry edge_sets`
-

-`info`
-
-`OriginInfo info`
-

-`node_sets`
-
-`repeated NodeSetsEntry node_sets`
-

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md
index 2c350198..819affc8 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md
@@ -1,50 +1,41 @@
 # tfgnn.GraphTensor
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A composite tensor for heterogeneous directed graphs with features.
 
 
 tfgnn.GraphTensor(
-    data: Data, spec: 'GraphPieceSpecBase', validate: bool = False
+    data: Data, spec: 'GraphPieceSpecBase'
 )
 

 
-
-
 
 
 A GraphTensor is an immutable container (as any composite tensor) to represent
 one or more heterogeneous directed graphs, as defined in the GraphTensor guide
 (or even hypergraphs). A GraphTensor consists of NodeSets, EdgeSets and a
-Context (collectively known as graph pieces), which are also composite
-tensors. The graph pieces consist of fields, which are `tf.Tensor`s and/or
-`tf.RaggedTensor`s that store the graph structure (esp. the edges between
-nodes) and user-defined features.
+Context (collectively known as graph pieces), which are also composite tensors.
+The graph pieces consist of fields, which are `tf.Tensor`s and/or
+`tf.RaggedTensor`s that store the graph structure (esp. the edges between nodes)
+and user-defined features.
 
-In the same way as `tf.Tensor` has numbers as its elements, the elements of
-the GraphTensor are graphs. Its `shape` of `[]` describes a scalar (single)
-graph, a shape of `[d0]` describes a `d0`-vector of graphs, a shape of
-`[d0, d1]` a `d0` x `d1` matrix of graphs, and so on.
+In the same way as `tf.Tensor` has numbers as its elements, the elements of the
+GraphTensor are graphs. Its `shape` of `[]` describes a scalar (single) graph, a
+shape of `[d0]` describes a `d0`-vector of graphs, a shape of `[d0, d1]` a `d0`
+x `d1` matrix of graphs, and so on.
 
 RULE: In the shape of a GraphTensor, no dimension except the outermost is
-allowed to be `None` (that is, of unknown size). Future versions of
-GraphTensor may lift that requirement.
+allowed to be `None` (that is, of unknown size). Future versions of GraphTensor
+may lift that requirement.
 
 Each of those graphs in a GraphTensor consists of 0 or more disjoint
-(sub-)graphs called graph components. The number of components could vary
-from graph to graph or be fixed to a value known statically. On a batched
+(sub-)graphs called graph components. The number of components could vary from
+graph to graph or be fixed to a value known statically. On a batched
 GraphTensor, one can call the method merge_batch_to_components() to merge all
 graphs of the batch into one, contiguously indexed graph containing the same
 components as the original graph tensor. See the GraphTensor guide for the
@@ -69,17 +60,18 @@ tfgnn.GraphTensor.from_pieces(
                 target=('node', [3, 7, 1])))})
 ```
 
-All graph pieces provide a mapping interface to access their features by name
-as `graph_piece[feature_name]`. Each graph piece feature has the shape
-`[*graph_shape, num_items, *feature_shape]`, where `graph_shape` is the shape
-of the GraphTensor, `num_items` is the number of items in a piece (number of
-graph components, number of nodes in a node set or edges in an edge set). The
+All graph pieces provide a mapping interface to access their features by name as
+`graph_piece[feature_name]`. Each graph piece feature has the shape
+`[*graph_shape, num_items, *feature_shape]`, where `graph_shape` is the shape of
+the GraphTensor, `num_items` is the number of items in a piece (number of graph
+components, number of nodes in a node set or edges in an edge set). The
 `feature_shape` is the shape of the feature value for each item.
 
-Naturally, the first GraphTensor.rank dimensions of all graph tensor fields
-must index the same graphs, the item dimension must correspond to the same
-item (graph component, node or edge) within the same graph piece (context,
-node set or edge set).
+Naturally, the first
+GraphTensor.rank
+dimensions of all graph tensor fields must index the same graphs, the item
+dimension must correspond to the same item (graph component, node or edge)
+within the same graph piece (context, node set or edge set).
 
 RULE: 'None' always denotes ragged or outermost field dimension. Uniform
 dimensions must have a fixed size that is given in the dimension.
@@ -113,8 +105,8 @@ venues = tfgnn.GraphTensor.from_pieces(
 ```
 
 The assignment of an item to its graph components is stored as the `sizes`
-attribute of the graph piece. Its shape is `[*graph_shape, num_components]`
-(the same for all pieces). The stored values are number of items in each graph
+attribute of the graph piece. Its shape is `[*graph_shape, num_components]` (the
+same for all pieces). The stored values are number of items in each graph
 component (so the `.sizes` has all 1s as its values).
 
 #### Example 3:
@@ -130,26 +122,26 @@ first_paper_year = tf.reduce_min(years_by_venue, -1)  # [2017, 2022]
 ```
 
 The GraphTensor, as a composite tensor, can be used directly in a
-tf.data.Dataset, as an input or output of a Keras Layer or a tf.function,
-and so on. As any other tensor, the GraphTensor has an associated type
-specification object, a GraphTensorSpec, that holds the `tf.TensorSpec` and
-`tf.RaggedTensorSpec` objects for all its fields, plus a bit of metadata such
-as graph connectivity (see tfgnn.GraphTensorSpec).
+tf.data.Dataset, as an input or output of a Keras Layer or a tf.function, and so
+on. As any other tensor, the GraphTensor has an associated type specification
+object, a GraphTensorSpec, that holds the `tf.TensorSpec` and
+`tf.RaggedTensorSpec` objects for all its fields, plus a bit of metadata such as
+graph connectivity (see tfgnn.GraphTensorSpec).
 
 The GraphTensor allows batching of graphs. Batching changes a GraphTensor
-instance's shape to `[batch_size, *graph_shape]` and the GraphTensor's
-rank is increased by 1. Unbatching removes dimension-0, as if truncating with
-`shape[1:]`, and the GraphTensor's rank is decreased by 1. This works
-naturally with the batch and unbatch methods of tf.data.Datatset.
+instance's shape to `[batch_size, *graph_shape]` and the GraphTensor's rank is
+increased by 1. Unbatching removes dimension-0, as if truncating with
+`shape[1:]`, and the GraphTensor's rank is decreased by 1. This works naturally
+with the batch and unbatch methods of tf.data.Datatset.
 
 RULE: Batching followed by unbatching results in a dataset with equal
-GraphTensors as before, except for the last incomplete batch (if batching
-used `drop_remainder=True`).
+GraphTensors as before, except for the last incomplete batch (if batching used
+`drop_remainder=True`).
 
 For now, GraphTensor requires that GraphTensor.shape does not contain `None`,
-except maybe as the outermost dimension. That means repeated calls to
-`.batch()` must set `drop_remainder=True` in all but the last one. Future
-versions of GraphTensor may lift that requirement.
+except maybe as the outermost dimension. That means repeated calls to `.batch()`
+must set `drop_remainder=True` in all but the last one. Future versions of
+GraphTensor may lift that requirement.
 
 All pieces and its fields are batched together with their GraphTensor so that
 shapes of a graph tensor, its pieces and features are all in sync.
@@ -159,67 +151,93 @@ ragged dimension of a RaggedTensor. (Note this is only allowed for the items
 dimension, not a graph dimension.) In all other cases, the type of the field
 (Tensor or RaggedTensor) is preserved.
 
+Graph tensor allows `int64` or `int32` types to index graph items. There are two
+types of indices: `indices_dtype` and `row_splits_dtype`. The `indices_dtype` is
+used to index itemes within graph pieces (`sizes`) and as a `dtype` of adjacency
+indices. The `row_splits_dtype` is a dtype for all ragged row partitions of all
+GraphTensor fields of type `tf.RaggedTensor`.
+
+RULE: `indices_dtype` and `row_splits_dtype` are consistent for all graph pieces
+within the graph tensor.
+
+IMPORTANT: This behaviour is disabled when loading legacy SavedModels created
+before this requirement was introduced. It is strongly recommented to align
+indices for all graph tensors generated by those legacy models using the methods
+`.with_indices_dtype()` and `.with_row_splits_dtype()`.
+
+The `indices_dtype` is `int32` by default, the default integer type in
+Tensorflow The `indices_dtype` for graph tensor and its pieces can be changed
+using `.with_indices_dtype()` method.
+
+The `row_splits_dtype` is `int64` by default, the same as for `RaggedTensor`s.
+They can be changed using `.with_row_splits_dtype()` method.
+
+NOTE: graph tensors can be constructed from pieces with inconsistent
+`indices_dtype` and `row_splits_dtype`. The indices types of the result
+`GraphTensor` are resolved towards the integer types with the maximum capacity
+and all pieces are casted towards those types. For example, if *any* graph piece
+used in `.from_pieces()` has `int64` `indices_dtype` the result graph tensor
+(and all its pieces) would have `int64` `indices_dtype`.
+
 
+
  
-
-
-

 
 
Args

 
 

 
-`data`
+data
 
 
 Nest of Field or subclasses of GraphPieceBase.
 
 

 
-`spec`
-
-A subclass of GraphPieceSpecBase with a `_data_spec` that matches
-`data`.
-

-`validate`
+spec
 
 
-if set, checks that data and spec are aligned, compatible and
-supported.
+A subclass of GraphPieceSpecBase with a _data_spec that matches
+data.
 
 

 

 
 
+
  
 
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout_into_feature.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout_into_feature.md
index 956f2a0c..583ef1cb 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout_into_feature.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout_into_feature.md
@@ -1,17 +1,10 @@
 # tfgnn.structured_readout_into_feature
 
-[TOC]
-
 
 
-

 
 
Attributes
  `context`   The graph context. 
  `edge_sets`   A read-only mapping
-from node set name to the node set.  
 
-`indices_dtype`   The integer type to
-represent ragged splits.  
  `node_sets`
-  A read-only mapping from node set name to the node set. 
  `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
+
  context   The graph context.
+ 
  edge_sets   A
+read-only mapping from node set name to the node set.  
 
+indices_dtype   The dtype for
+graph items indexing. One of tf.int32 or tf.int64.
+ 
  node_sets   A
+read-only mapping from node set name to the node set.  
 
+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 TensorShape is guaranteed to have a known rank, but the
 individual dimension sizes may be unknown.
 
 

 
-`spec`
+spec
 
 
 The public type specification of this tensor.
 
 

 
-`total_num_components`
+total_num_components
 
 
 The total number of graph components.
@@ -231,7 +249,7 @@ The total number of graph components.
 
 
from_pieces

 
-View
+View
 source
 
 
@@ -239,16 +257,16 @@ source
 from_pieces(
     context: Optional[Context] = None,
     node_sets: Optional[Mapping[NodeSetName, NodeSet]] = None,
-    edge_sets: Optional[Mapping[EdgeSetName, EdgeSet]] = None
+    edge_sets: Optional[Mapping[EdgeSetName, EdgeSet]] = None,
+    validate: Optional[bool] = None
 ) -> 'GraphTensor'
 

 
 Constructs a new `GraphTensor` from context, node sets and edge sets.
 
-
 
merge_batch_to_components

 
-View
+View
 source
 
 
@@ -257,13 +275,12 @@ source
 
 Merges all contained graphs into one contiguously indexed graph.
 
-On a batched GraphTensor, one can call this method merge all graphs of the
-batch into one, contiguously indexed graph. The resulting GraphTensor has
-shape [] (i.e., is scalar) and its features have the shape
-`[total_num_items, *feature_shape]` where `total_num_items` is the sum of
-the previous `num_items` per batch element. Most TF-GNN models expect scalar
-GraphTensors. There is no function to reverse this method.
-
+On a batched GraphTensor, one can call this method merge all graphs of the batch
+into one, contiguously indexed graph. The resulting GraphTensor has shape `[]`
+(i.e., is scalar) and its features have the shape `[total_num_items,
+*feature_shape]` where `total_num_items` is the sum of the previous `num_items`
+per batch element. Most TF-GNN models expect scalar GraphTensors. There is no
+function to reverse this method.
 
 Example: Flattening of
 
@@ -307,13 +324,14 @@ tfgnn.GraphTensor.from_pieces(
             sizes=[3, 2, 1, 1],
             features={},
             # Note how node indices have changes to reference nodes
-            # withing the same graph ignoring its components.
+            # within the same graph ignoring its components.
             adjacency=tfgnn.Adjacency.from_indices(
                 source=('node', [0, 1, 1, 2, 3 + 0, 3 + 1 + 0]),
                 target=('node', [0, 0, 1, 2, 3 + 0, 3 + 1 + 0])))})
 ```
 
 
+
  
@@ -325,11 +343,9 @@ A scalar (rank 0) graph tensor.
 
 

 
 
Returns

 
-
-
 
remove_features

 
-View
+View
 source
 
 
@@ -360,8 +376,6 @@ result = graph.remove_features(node_sets={'node.a': ['id']})
 
 #### Result:
 
-
-
 ```python
 tfgnn.GraphTensor.from_pieces(
     node_sets={
@@ -374,20 +388,21 @@ tfgnn.GraphTensor.from_pieces(
 ```
 
 
+
  

 
 
Args

 
 

 
-`context`
+context
 
 
-A list of feature names to remove from the context, or `None`.
+A list of feature names to remove from the context, or None.
 
 

 
-`node_sets`
+node_sets
 
 
 A mapping from node set names to lists of feature names to be
@@ -395,7 +410,7 @@ removed from the respective node sets.
 
 

 
-`edge_sets`
+edge_sets
 
 
 A mapping from edge set names to lists of feature names to be
@@ -404,15 +419,14 @@ removed from the respective edge sets.
 

 

 
-
-
 
+
  
@@ -420,16 +434,15 @@ to be removed or is still present in the output.
 
 

 
 
Returns

 

 
-A `GraphTensor` with the same graph topology as the input and a subset
+A GraphTensor with the same graph topology as the input and a subset
 of its features. Each feature of the input either was named as a feature
 to be removed or is still present in the output.
 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if some feature names in the arguments were not present in the
@@ -438,11 +451,9 @@ input graph tensor.
 

 

 
-
-
 
replace_features

 
-View
+View
 source
 
 
@@ -478,8 +489,6 @@ result = graph.replace_features(
 
 #### Result:
 
-
-
 ```python
 tfgnn.GraphTensor.from_pieces(
     context=tfgnn.Context.from_fields(
@@ -499,22 +508,23 @@ tfgnn.GraphTensor.from_pieces(
 ```
 
 
+
  

 
 
Args

 
 

 
-`context`
+context
 
 
-A substitute for the context features, or `None` (which keeps the
+A substitute for the context features, or None (which keeps the
 prior features). Their tensor shapes must match the number of existing
 components, which remains unchanged.
 
 

 
-`node_sets`
+node_sets
 
 
 Substitutes for the features of the specified node sets. Their
@@ -523,7 +533,7 @@ unchanged. Features on node sets that are not included remain unchanged.
 
 

 
-`edge_sets`
+edge_sets
 
 
 Substitutes for the features of the specified edge sets. Their
@@ -534,31 +544,29 @@ are not included remain unchanged.
 

 

 
-
-
 
+
  

 
 
Returns

 

 
-A `GraphTensor` instance with feature maps replaced according to the
+A GraphTensor instance with feature maps replaced according to the
 arguments.
 
 

 
 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if some node sets or edge sets are not present in the graph
@@ -567,21 +575,54 @@ tensor.
 

 

 
-
-
 
set_shape

 
-View
+View
 source
 
 
 set_shape(
     new_shape: ShapeLike
-) -> 'GraphPieceSpecBase'
+) -> 'GraphPieceBase'
 

 
-Enforce the common prefix shape on all the contained features.
+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.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md
index 623fba59..2fadd307 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md
@@ -1,59 +1,54 @@
 # tfgnn.GraphTensorSpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
-A type spec for tfgnn.GraphTensor.
+A type spec for
+tfgnn.GraphTensor.
 
 
 tfgnn.GraphTensorSpec(
     data_spec: DataSpec,
     shape: tf.TensorShape,
     indices_dtype: tf.dtypes.DType,
+    row_splits_dtype: tf.dtypes.DType,
     metadata: Metadata = None,
-    validate: bool = False
+    check_consistent_indices_dtype: bool = False,
+    check_consistent_row_splits_dtype: bool = False
 )
 

 
-
-
 
-
-
-
-
 
+
  
 
-

 
 
Attributes
  `context_spec`   The graph context
-type spec.  
  `edge_sets_spec`
-  A read-only mapping form edge set name to the edge set type spec.
- 
  `indices_dtype`   The
-integer type to represent ragged splits.  
 
-`node_sets_spec`   A read-only mapping form
-node set name to the node set type spec.  
 
-`rank`   The rank of the GraphPiece. Guaranteed not to
-be `None`.  
  `shape`   A
-possibly-partial shape specification of the GraphPiece.
-
-The returned `TensorShape` is guaranteed to have a known rank, but the
-individual dimension sizes may be unknown.  
 
-`total_num_components`   The total
-number of graph components if known.  
 
-`value_type`   The Python type for values that
-are compatible with this TypeSpec.
+
  context_spec   The
+graph context type spec.  
 
+edge_sets_spec   A read-only
+mapping form edge set name to the edge set type spec.  
 
+indices_dtype   The dtype for
+graph items indexing. One of tf.int32 or tf.int64.
+ 
  node_sets_spec
+  A read-only mapping form node set name to the node set type spec.
+ 
  rank   The rank of
+the GraphPiece. 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 of the GraphPiece.
+
+The returned TensorShape is guaranteed to have a known rank, but
+the individual dimension sizes may be unknown.  
 
+total_num_components  
+The total number of graph components if known.  
 
+value_type   The Python type for
+values that are compatible with this TypeSpec.
 
 In particular, all values that are compatible with this TypeSpec must be an
 instance of this type.
@@ -61,8 +56,6 @@ instance of this type.
 

 

 
-
-
 ## Methods
 
 
experimental_as_proto

@@ -89,13 +82,14 @@ Returns a TypeSpec instance based on the serialized proto.
 Do NOT override for custom non-TF types.
 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/satisfies_size_constraints.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/satisfies_size_constraints.md
index 55d29832..9c4ad426 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/satisfies_size_constraints.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/satisfies_size_constraints.md
@@ -1,17 +1,10 @@
 # tfgnn.satisfies_size_constraints
 
-[TOC]
-
 
 
-

 
 
Args

 
 

 
-`proto`
+proto
 
 
 Proto generated using 'experimental_as_proto'.
@@ -116,7 +110,7 @@ Do NOT override for custom non-TF types.
 
 
from_piece_specs

 
-View
+View
 source
 
 
@@ -128,12 +122,13 @@ source
 ) -> 'GraphTensorSpec'
 

 
-The counterpart of GraphTensor.from_pieces for pieces type specs.
-
+The counterpart of
+GraphTensor.from_pieces
+for pieces type specs.
 
 
from_value

 
-View
+View
 source
 
 
@@ -145,7 +140,6 @@ source
 
 Extension Types API: Factory method.
 
-
 
is_compatible_with

 
 
@@ -160,13 +154,14 @@ Prefer using "is_subtype_of" and "most_specific_common_supertype" wherever
 possible.
 
 
+
  

 
 
Args

 
 

 
-`spec_or_value`
+spec_or_value
 
 
 A TypeSpec or TypeSpec associated value to compare against.
@@ -174,8 +169,6 @@ A TypeSpec or TypeSpec associated value to compare against.
 

 

 
-
-
 
is_subtype_of

 
 
@@ -188,19 +181,19 @@ Returns True if `self` is a subtype of `other`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
 A TraceType object.
@@ -208,8 +201,6 @@ A TraceType object.
 

 

 
-
-
 
most_specific_common_supertype

 
 
@@ -218,23 +209,23 @@ A TraceType object.
 ) -> Optional['TypeSpec']
 

 
-Returns the most specific supertype TypeSpec  of `self` and `others`.
+Returns the most specific supertype TypeSpec of `self` and `others`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`others`
+others
 
 
 A sequence of TraceTypes.
@@ -242,8 +233,6 @@ A sequence of TraceTypes.
 

 

 
-
-
 
most_specific_compatible_type

 
 
@@ -252,51 +241,51 @@ A sequence of TraceTypes.
 ) -> 'TypeSpec'
 

 
-Returns the most specific TypeSpec compatible with `self` and `other`. (deprecated)
+Returns the most specific TypeSpec compatible with `self` and `other`.
+(deprecated)
 
 Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version.
-Instructions for updating:
-Use most_specific_common_supertype instead.
+Instructions for updating: Use most_specific_common_supertype instead.
 
-Deprecated. Please use `most_specific_common_supertype` instead.
-Do not override this function.
+Deprecated. Please use `most_specific_common_supertype` instead. Do not override
+this function.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
-A `TypeSpec`.
+A TypeSpec.
 
 

 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If there is no TypeSpec that is compatible with both `self`
-and `other`.
+If there is no TypeSpec that is compatible with both self
+and other.
 
 

 

 
 
relax

 
-View
+View
 source
 
 
@@ -313,27 +302,28 @@ Allows variable number of graph nodes, edges or/and graph components.
 Calling with all default parameters keeps the spec unchanged.
 
 
+
  

 
 
Args

 
 

 
-`num_components`
+num_components
 
 
 if True, allows a variable number of graph components.
 
 

 
-`num_nodes`
+num_nodes
 
 
 if True, allows a variable number of nodes in each node set.
 
 

 
-`num_edges`
+num_edges
 
 
 if True, allows a variable number of edges in each edge set.
@@ -342,6 +332,7 @@ if True, allows a variable number of edges in each edge set.
 

 
 
+
  
@@ -354,13 +345,14 @@ Relaxed compatible graph tensor spec.
 

 
 
Returns

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if graph tensor is not scalar (rank > 0).
@@ -368,6 +360,45 @@ if graph tensor is not scalar (rank > 0).
 

 

 
+
with_indices_dtype

+
+View
+source
+
+
+with_indices_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_row_splits_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_shape(
+    new_shape: ShapeLike
+) -> 'GraphPieceSpecBase'
+

+
+Enforce the common prefix shape on all the contained features.
+
 
__eq__

 
 
@@ -378,7 +409,6 @@ if graph tensor is not scalar (rank > 0).
 
 Return self==value.
 
-
 
__ne__

 
 
@@ -388,7 +418,3 @@ Return self==value.
 

 
 Return self!=value.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacency.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacency.md
index f5df0954..98b2aadd 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacency.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacency.md
@@ -1,93 +1,82 @@
 # tfgnn.HyperAdjacency
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Stores how (hyper-)edges connect tuples of nodes from incident node sets.
 
 
 tfgnn.HyperAdjacency(
-    data: Data, spec: 'GraphPieceSpecBase', validate: bool = False
+    data: Data, spec: 'GraphPieceSpecBase'
 )
 

 
-
-
 
 
-The incident node sets in the hyper-adjacency are referenced by a unique
-integer identifier called the node set tag. (For a non-hyper graph, it is
-conventional to use the integers tfgnn.SOURCE and tfgnn.TARGET.) This
-allows the hyper-adjacency to connect nodes from the same or different node
-sets. Each hyper-edge connects a fixed number of nodes, one node from each
-incident node set. The adjacency information is stored as a mapping from the
-node set tags to integer tensors containing indices of nodes in corresponding
-node sets. Those tensors are indexed by edges. All index tensors have the same
-type spec and shape of `[*graph_shape, num_edges]`, where `num_edges` is the
-number of edges in the edge set (could be potentially ragged). The index
-tensors are of `tf.Tensor` type if `num_edges` is not `None` or
-`graph_shape.rank = 0` and of `tf.RaggedTensor` type otherwise.
+The incident node sets in the hyper-adjacency are referenced by a unique integer
+identifier called the node set tag. (For a non-hyper graph, it is conventional
+to use the integers tfgnn.SOURCE
+and tfgnn.TARGET.) This allows the
+hyper-adjacency to connect nodes from the same or different node sets. Each
+hyper-edge connects a fixed number of nodes, one node from each incident node
+set. The adjacency information is stored as a mapping from the node set tags to
+integer tensors containing indices of nodes in corresponding node sets. Those
+tensors are indexed by edges. All index tensors have the same type spec and
+shape of `[*graph_shape, num_edges]`, where `num_edges` is the number of edges
+in the edge set (could be potentially ragged). The index tensors are of
+`tf.Tensor` type if `num_edges` is not `None` or `graph_shape.rank = 0` and of
+`tf.RaggedTensor` type otherwise.
 
 The HyperAdjacency is a composite tensor.
 
 
+
  
-
-
-

 
 
Args

 
 

 
-`data`
+data
 
 
 Nest of Field or subclasses of GraphPieceBase.
 
 

 
-`spec`
+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.
+A subclass of GraphPieceSpecBase with a _data_spec that matches
+data.
 
 

 

 
 
+
  
 
-
-
@@ -128,7 +121,7 @@ given to constructor.
 
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/make_sampling_spec_tree.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/make_sampling_spec_tree.md
index 1730a42a..3b23d48b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/make_sampling_spec_tree.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/make_sampling_spec_tree.md
@@ -1,23 +1,16 @@
 # tfgnn.sampler.make_sampling_spec_tree
 
-[TOC]
-
 
 
-

 
 
Attributes
  `indices_dtype`   The integer type
-to represent ragged splits.  
  `rank`  The rank of this Tensor. Guaranteed not to be `None`.  
 
-`shape`   A possibly-partial shape specification for
-this Tensor.
+
  indices_dtype   The
+dtype for graph items indexing. One of tf.int32 or
+tf.int64.  
  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 `TensorShape` is guaranteed to have a known rank, but the
+The returned TensorShape is guaranteed to have a known rank, but the
 individual dimension sizes may be unknown.
 
 

 
-`spec`
+spec
 
 
 The public type specification of this tensor.
@@ -99,23 +88,20 @@ The public type specification of this tensor.
 
 
from_indices

 
-View
+View
 source
 
 
 @classmethod
 from_indices(
-    indices: Indices, *_, validate: bool = True
+    indices: Indices, *_, validate: Optional[bool] = None
 ) -> 'HyperAdjacency'
 

 
 Constructs a new instance from the `indices` tensors.
 
-
 #### Example 1:
 
-
-
 ```python
 # Single graph (rank is 0). Connects pairs of nodes (a[0], b[2]),
 # (a[1], b[1]), (a[2], b[0]) from node sets a and b.
@@ -127,8 +113,6 @@ tfgnn.HyperAdjacency.from_indices({
 
 #### Example 2:
 
-
-
 ```python
 # Single hypergraph (rank is 0). Connects triplets of nodes
 # (a[0], b[2], c[1]), (a[1], b[1], c[0]) from the node sets a, b and c.
@@ -141,8 +125,6 @@ tfgnn.HyperAdjacency.from_indices({
 
 #### Example 3:
 
-
-
 ```python
 # Batch of two graphs (rank is 1). Connects pairs of nodes in
 # graph 0: (a[0], b[2]), (a[1], b[1]); graph 1: (a[2], b[0]).
@@ -153,52 +135,50 @@ tfgnn.HyperAdjacency.from_indices({
 ```
 
 
+
  

 
 
Args

 
 

 
-`indices`
+indices
 
 
 A mapping from node tags to 2-tuples of node set name and node
 index tensor. The index tensors must have the same type spec and shape
-of `[*graph_shape, num_edges]`, where `num_edges` is the number of edges
-in each graph (could be ragged). The index tensors are of `tf.Tensor`
-type if `num_edges` is not `None` or `graph_shape.rank = 0` and of
-`tf.RaggedTensor` type otherwise.
+of [*graph_shape, num_edges], where num_edges is the number of edges
+in each graph (could be ragged). The index tensors are of tf.Tensor
+type if num_edges is not None or graph_shape.rank = 0 and of
+tf.RaggedTensor type otherwise.
 
 

 
-`validate`
+validate
 
 
-If `True`, checks that node indices have the same type spec.
+If True, checks that node indices have the same type spec.
 
 

 

 
-
-
 
+
  

 
 
Returns

 

 
-A `HyperAdjacency` tensor with its `shape` and `indices_dtype` being
-inferred from the passed `indices` values.
+A HyperAdjacency tensor with its shape and indices_dtype being
+inferred from the passed indices values.
 
 

 
 

 
-
-
 
get_indices_dict

 
-View
+View
 source
 
 
@@ -207,10 +187,9 @@ source
 
 Returns copy of indices as a dictionary.
 
-
 
node_set_name

 
-View
+View
 source
 
 
@@ -221,24 +200,61 @@ source
 
 Returns a node set name for the given node set tag.
 
-
 
set_shape

 
-View
+View
 source
 
 
 set_shape(
     new_shape: ShapeLike
-) -> 'GraphPieceSpecBase'
+) -> 'GraphPieceBase'
 

 
-Enforce the common prefix shape on all the contained features.
+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
+View
 source
 
 
@@ -248,7 +264,3 @@ source
 

 
 Returns an index tensor for the given node set tag.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacencySpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacencySpec.md
index 25df52d3..a25fe62e 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacencySpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/HyperAdjacencySpec.md
@@ -1,53 +1,48 @@
 # tfgnn.HyperAdjacencySpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
-A type spec for tfgnn.HyperAdjacency.
+A type spec for
+tfgnn.HyperAdjacency.
 
 
 tfgnn.HyperAdjacencySpec(
     data_spec: DataSpec,
     shape: tf.TensorShape,
     indices_dtype: tf.dtypes.DType,
+    row_splits_dtype: tf.dtypes.DType,
     metadata: Metadata = None,
-    validate: bool = False
+    check_consistent_indices_dtype: bool = False,
+    check_consistent_row_splits_dtype: bool = False
 )
 

 
-
-
 
-
-
-
-
 
+
  
 
-
-
-

 
 
Attributes
  `indices_dtype`   The integer type
-to represent ragged splits.  
  `rank`  The rank of the GraphPiece. Guaranteed not to be `None`.  
 `shape`   A possibly-partial shape specification
-of the GraphPiece.
-
-The returned `TensorShape` is guaranteed to have a known rank, but the
-individual dimension sizes may be unknown.  
 
-`total_size`   The total number of edges if
-known.  
  `value_type`   The
-Python type for values that are compatible with this TypeSpec.
+
  indices_dtype   The
+dtype for graph items indexing. One of tf.int32 or
+tf.int64.  
  rank
+  The rank of the GraphPiece. 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 of the GraphPiece.
+
+The returned TensorShape is guaranteed to have a known rank, but
+the individual dimension sizes may be unknown.  
 
+total_size   The total number of
+edges if known.  
 
+value_type   The Python type for
+values that are compatible with this TypeSpec.
 
 In particular, all values that are compatible with this TypeSpec must be an
 instance of this type.
@@ -55,8 +50,6 @@ instance of this type.
 

 

 
-
-
 ## Methods
 
 
experimental_as_proto

@@ -83,13 +76,14 @@ Returns a TypeSpec instance based on the serialized proto.
 Do NOT override for custom non-TF types.
 
 
+
  
@@ -85,7 +78,7 @@ its size in each graph component must be 1.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddSelfLoops.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddSelfLoops.md
index 04bfac0b..69dab1e2 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddSelfLoops.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddSelfLoops.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.AddSelfLoops
 
-[TOC]
-
 
 
-

 
 
Args

 
 

 
-`proto`
+proto
 
 
 Proto generated using 'experimental_as_proto'.
@@ -110,7 +104,7 @@ Do NOT override for custom non-TF types.
 
 
from_incident_node_sets

 
-View
+View
 source
 
 
@@ -123,52 +117,49 @@ source
 
 Constructs a new instance from the `incident_node_sets`.
 
-
 
+
  

 
 
Args

 
 

 
-`incident_node_sets`
+incident_node_sets
 
 
 A mapping from incident node tags to node set names.
 
 

 
-`index_spec`
+index_spec
 
 
 type spec for all index tensors of shape
-`[*graph_shape, num_edges]`, where `num_edges` is the number of edges in
-each graph. If `num_edges` is not `None` or `graph_shape.rank = 0` the
-spec must be of `tf.TensorSpec` type and of `tf.RaggedTensorSpec` type
+[*graph_shape, num_edges], where num_edges is the number of edges in
+each graph. If num_edges is not None or graph_shape.rank = 0 the
+spec must be of tf.TensorSpec type and of tf.RaggedTensorSpec type
 otherwise.
 
 

 

 
-
-
 
+
  

 
 
Returns

 

 
-A `HyperAdjacencySpec` TypeSpec.
+A HyperAdjacencySpec TypeSpec.
 
 

 
 

 
-
-
 
from_value

 
-View
+View
 source
 
 
@@ -180,10 +171,9 @@ source
 
 Extension Types API: Factory method.
 
-
 
get_index_specs_dict

 
-View
+View
 source
 
 
@@ -192,7 +182,6 @@ source
 
 Returns copy of indices type specs as a dictionary.
 
-
 
is_compatible_with

 
 
@@ -207,13 +196,14 @@ Prefer using "is_subtype_of" and "most_specific_common_supertype" wherever
 possible.
 
 
+
  

 
 
Args

 
 

 
-`spec_or_value`
+spec_or_value
 
 
 A TypeSpec or TypeSpec associated value to compare against.
@@ -221,8 +211,6 @@ A TypeSpec or TypeSpec associated value to compare against.
 

 

 
-
-
 
is_subtype_of

 
 
@@ -235,19 +223,19 @@ Returns True if `self` is a subtype of `other`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
 A TraceType object.
@@ -255,8 +243,6 @@ A TraceType object.
 

 

 
-
-
 
most_specific_common_supertype

 
 
@@ -265,23 +251,23 @@ A TraceType object.
 ) -> Optional['TypeSpec']
 

 
-Returns the most specific supertype TypeSpec  of `self` and `others`.
+Returns the most specific supertype TypeSpec of `self` and `others`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`others`
+others
 
 
 A sequence of TraceTypes.
@@ -289,8 +275,6 @@ A sequence of TraceTypes.
 

 

 
-
-
 
most_specific_compatible_type

 
 
@@ -299,53 +283,51 @@ A sequence of TraceTypes.
 ) -> 'TypeSpec'
 

 
-Returns the most specific TypeSpec compatible with `self` and `other`. (deprecated)
+Returns the most specific TypeSpec compatible with `self` and `other`.
+(deprecated)
 
 Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version.
-Instructions for updating:
-Use most_specific_common_supertype instead.
+Instructions for updating: Use most_specific_common_supertype instead.
 
-Deprecated. Please use `most_specific_common_supertype` instead.
-Do not override this function.
+Deprecated. Please use `most_specific_common_supertype` instead. Do not override
+this function.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
-A `TypeSpec`.
+A TypeSpec.
 
 

 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If there is no TypeSpec that is compatible with both `self`
-and `other`.
+If there is no TypeSpec that is compatible with both self
+and other.
 
 

 

 
-
-
 
node_set_name

 
-View
+View
 source
 
 
@@ -358,7 +340,7 @@ Returns a node set name for the given node set tag.
 
 
relax

 
-View
+View
 source
 
 
@@ -372,13 +354,14 @@ Allows variable number of graph edges.
 Calling with all default parameters keeps the spec unchanged.
 
 
+
  

 
 
Args

 
 

 
-`num_edges`
+num_edges
 
 
 if True, allows a variable number of edges in each edge set.
@@ -388,6 +371,7 @@ If False, returns spec unchanged.
 

 
 
+
  
@@ -400,13 +384,14 @@ Relaxed compatible spec.
 

 
 
Returns

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if adjacency is not scalar (rank > 0).
@@ -414,6 +399,45 @@ if adjacency is not scalar (rank > 0).
 

 

 
+
with_indices_dtype

+
+View
+source
+
+
+with_indices_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_row_splits_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_shape(
+    new_shape: ShapeLike
+) -> 'GraphPieceSpecBase'
+

+
+Enforce the common prefix shape on all the contained features.
+
 
__eq__

 
 
@@ -424,10 +448,9 @@ if adjacency is not scalar (rank > 0).
 
 Return self==value.
 
-
 
__getitem__

 
-View
+View
 source
 
 
@@ -438,7 +461,6 @@ source
 
 Returns an index tensor type spec for the given node set tag.
 
-
 
__ne__

 
 
@@ -448,7 +470,3 @@ Returns an index tensor type spec for the given node set tag.
 

 
 Return self!=value.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/IncidentNodeOrContextTag.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/IncidentNodeOrContextTag.md
index 4fd13809..4bacf66d 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/IncidentNodeOrContextTag.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/IncidentNodeOrContextTag.md
@@ -1,14 +1,8 @@
-

-
-
-

-
 # tfgnn.IncidentNodeOrContextTag
 
 
-This symbol is a **type alias**.
-
 
+This symbol is a **type alias**.
 
 #### Source:
 
@@ -19,6 +13,4 @@ This symbol is a **type alias**.
 ]
 

 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md
index 2ad4b413..2d56c857 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md
@@ -1,34 +1,25 @@
 # tfgnn.NodeSet
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A composite tensor for node set features plus size information.
 
 
 tfgnn.NodeSet(
-    data: Data, spec: 'GraphPieceSpecBase', validate: bool = False
+    data: Data, spec: 'GraphPieceSpecBase'
 )
 

 
-
-
 
 
 The items of the node set are subset of graph nodes.
 
-All nodes in a node set have the same features, identified by a string key.
-Each feature is stored as one tensor and has shape `[*graph_shape, num_nodes,
+All nodes in a node set have the same features, identified by a string key. Each
+feature is stored as one tensor and has shape `[*graph_shape, num_nodes,
 *feature_shape]`. The `num_nodes` is the number of nodes in a graph (could be
 ragged). The `feature_shape` is the shape of the feature value for each node.
 NodeSet supports both fixed-size and variable-size features. The fixed-size
@@ -36,83 +27,80 @@ features must have fully defined feature_shape. They are stored as `tf.Tensor`
 if `num_nodes` is fixed-size or `graph_shape.rank = 0`. Variable-size node
 features are always stored as `tf.RaggedTensor`.
 
-Note that node set features are indexed without regard to graph components.
-The information which node belong to which graph component is contained in
-the `.sizes` tensor which defines the number of nodes in each graph component.
+Note that node set features are indexed without regard to graph components. The
+information which node belong to which graph component is contained in the
+`.sizes` tensor which defines the number of nodes in each graph component.
 
 
+
  
-
-
-

 
 
Args

 
 

 
-`data`
+data
 
 
 Nest of Field or subclasses of GraphPieceBase.
 
 

 
-`spec`
+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.
+A subclass of GraphPieceSpecBase with a _data_spec that matches
+data.
 
 

 

 
 
+
  
 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers.md
index 82cf3f2c..0c3b0049 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers.md
@@ -1,24 +1,13 @@
 # Module: tfgnn.keras.layers
 
-[TOC]
-
 
 
-

 
 
Attributes
  `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
+
  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 TensorShape is guaranteed to have a known rank, but the
 individual dimension sizes may be unknown.
 
 

 
-`sizes`
+sizes
 
 
 The number of items in each graph component.
 
 

 
-`spec`
+spec
 
 
 The public type specification of this tensor.
 
 

 
-`total_num_components`
+total_num_components
 
 
 The total number of graph components.
 
 

 
-`total_size`
+total_size
 
 
 The total number of items.
@@ -124,23 +112,23 @@ The total number of items.
 
 
from_fields

 
-View
+View
 source
 
 
 @classmethod
 from_fields(
-    *, features: Optional[Fields] = None, sizes: Field
+    *_,
+    features: Optional[Fields] = None,
+    sizes: Field,
+    validate: Optional[bool] = None
 ) -> 'NodeSet'
 

 
 Constructs a new instance from node set fields.
 
-
 #### Example:
 
-
-
 ```python
 tfgnn.NodeSet.from_fields(
     sizes=tf.constant([3]),
@@ -155,52 +143,60 @@ tfgnn.NodeSet.from_fields(
 ```
 
 
+
  
+
+
+

 
 
Args

 
 

 
-`features`
+features
 
 
 A mapping from feature name to feature Tensors or RaggedTensors.
-All feature tensors must have shape `[*graph_shape, num_nodes,
-*feature_shape]`, where `num_nodes` is the number of nodes in the node
+All feature tensors must have shape [*graph_shape, num_nodes,
+*feature_shape], where num_nodes is the number of nodes in the node
 set (could be ragged) and feature_shape is a shape of the feature value
 for each node.
 
 

 
-`sizes`
+sizes
 
 
 A number of nodes in each graph component. Has shape
-`[*graph_shape, num_components]`, where `num_components` is the number
+[*graph_shape, num_components], where num_components is the number
 of graph components (could be ragged).
 

+validate
+
+If true, use tf.assert ops to inspect the shapes of each field
+and check at runtime that they form a valid NodeSet.  The default
+behavior is set by the disable_graph_tensor_validation_at_runtime()
+and enable_graph_tensor_validation_at_runtime().
+
 

 

 
-
-
 
+
  

 
 
Returns

 

 
-A `NodeSet` composite tensor.
+A NodeSet composite tensor.
 
 

 
 

 
-
-
 
get_features_dict

 
-View
+View
 source
 
 
@@ -209,10 +205,9 @@ source
 
 Returns features copy as a dictionary.
 
-
 
replace_features

 
-View
+View
 source
 
 
@@ -223,20 +218,57 @@ source
 
 Returns a new instance with a new set of features.
 
-
 
set_shape

 
-View
+View
 source
 
 
 set_shape(
     new_shape: ShapeLike
-) -> 'GraphPieceSpecBase'
+) -> 'GraphPieceBase'
 

 
-Enforce the common prefix shape on all the contained features.
+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__

 
@@ -250,7 +282,3 @@ source
 

 
 Indexing operator `[]` to access feature values by their name.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSetSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSetSpec.md
index f1cd9522..1742a1ce 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSetSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSetSpec.md
@@ -1,17 +1,10 @@
 # tfgnn.NodeSetSpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A type spec for tfgnn.NodeSet.
 
@@ -20,40 +13,41 @@ A type spec for tfgnn.NodeSet.
     data_spec: DataSpec,
     shape: tf.TensorShape,
     indices_dtype: tf.dtypes.DType,
+    row_splits_dtype: tf.dtypes.DType,
     metadata: Metadata = None,
-    validate: bool = False
+    check_consistent_indices_dtype: bool = False,
+    check_consistent_row_splits_dtype: bool = False
 )
 

 
-
-
 
-
-
-
-
 
+
  
 
-

 
 
Attributes
  `features_spec`   A read-only
-mapping of feature name to feature spec.  
 
-`indices_dtype`   The integer type to
-represent ragged splits.  
  `rank`  
-The rank of the GraphPiece. Guaranteed not to be `None`.  
 
-`shape`   A possibly-partial shape specification of
-the GraphPiece.
-
-The returned `TensorShape` is guaranteed to have a known rank, but the
-individual dimension sizes may be unknown.  
 
-`sizes_spec`   The type spec for the sizes that
-provides num. elements per component.  
 
-`total_num_components`   The total
-number of graph components if known.  
 
-`total_size`   The total number of graph items
-if known.  
  `value_type`  
-The Python type for values that are compatible with this TypeSpec.
+
  features_spec   A
+read-only mapping of feature name to feature spec.  
 
+indices_dtype   The dtype for
+graph items indexing. One of tf.int32 or tf.int64.
+ 
  rank   The rank of
+the GraphPiece. 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 of the GraphPiece.
+
+The returned TensorShape is guaranteed to have a known rank, but
+the individual dimension sizes may be unknown.  
 
+sizes_spec   The type spec for the
+sizes that provides num. elements per component.  
 
+total_num_components  
+The total number of graph components if known.  
 
+total_size   The total number of
+graph items if known.  
 
+value_type   The Python type for
+values that are compatible with this TypeSpec.
 
 In particular, all values that are compatible with this TypeSpec must be an
 instance of this type.
@@ -61,8 +55,6 @@ instance of this type.
 

 

 
-
-
 ## Methods
 
 
experimental_as_proto

@@ -89,13 +81,14 @@ Returns a TypeSpec instance based on the serialized proto.
 Do NOT override for custom non-TF types.
 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_io_spec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_io_spec.md
index a61de5d7..1092f2ba 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_io_spec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_io_spec.md
@@ -1,17 +1,10 @@
 # tfgnn.get_io_spec
 
-[TOC]
-
 
 
-

 
 
Args

 
 

 
-`proto`
+proto
 
 
 Proto generated using 'experimental_as_proto'.
@@ -116,7 +109,7 @@ Do NOT override for custom non-TF types.
 
 
from_field_specs

 
-View
+View
 source
 
 
@@ -126,12 +119,13 @@ source
 ) -> 'NodeSetSpec'
 

 
-The counterpart of NodeSet.from_fields() for values type specs.
-
+The counterpart of
+NodeSet.from_fields()
+for values type specs.
 
 
from_value

 
-View
+View
 source
 
 
@@ -143,7 +137,6 @@ source
 
 Extension Types API: Factory method.
 
-
 
is_compatible_with

 
 
@@ -158,13 +151,14 @@ Prefer using "is_subtype_of" and "most_specific_common_supertype" wherever
 possible.
 
 
+
  

 
 
Args

 
 

 
-`spec_or_value`
+spec_or_value
 
 
 A TypeSpec or TypeSpec associated value to compare against.
@@ -172,8 +166,6 @@ A TypeSpec or TypeSpec associated value to compare against.
 

 

 
-
-
 
is_subtype_of

 
 
@@ -186,19 +178,19 @@ Returns True if `self` is a subtype of `other`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
 A TraceType object.
@@ -206,8 +198,6 @@ A TraceType object.
 

 

 
-
-
 
most_specific_common_supertype

 
 
@@ -216,23 +206,23 @@ A TraceType object.
 ) -> Optional['TypeSpec']
 

 
-Returns the most specific supertype TypeSpec  of `self` and `others`.
+Returns the most specific supertype TypeSpec of `self` and `others`.
 
 Implements the tf.types.experimental.func.TraceType interface.
 
-If not overridden by a subclass, the default behavior is to assume the
-TypeSpec is covariant upon attributes that implement TraceType and
-invariant upon rest of the attributes as well as the structure and type
-of the TypeSpec.
+If not overridden by a subclass, the default behavior is to assume the TypeSpec
+is covariant upon attributes that implement TraceType and invariant upon rest of
+the attributes as well as the structure and type of the TypeSpec.
 
 
+
  

 
 
Args

 
 

 
-`others`
+others
 
 
 A sequence of TraceTypes.
@@ -240,8 +230,6 @@ A sequence of TraceTypes.
 

 

 
-
-
 
most_specific_compatible_type

 
 
@@ -250,51 +238,51 @@ A sequence of TraceTypes.
 ) -> 'TypeSpec'
 

 
-Returns the most specific TypeSpec compatible with `self` and `other`. (deprecated)
+Returns the most specific TypeSpec compatible with `self` and `other`.
+(deprecated)
 
 Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version.
-Instructions for updating:
-Use most_specific_common_supertype instead.
+Instructions for updating: Use most_specific_common_supertype instead.
 
-Deprecated. Please use `most_specific_common_supertype` instead.
-Do not override this function.
+Deprecated. Please use `most_specific_common_supertype` instead. Do not override
+this function.
 
 
+
  

 
 
Args

 
 

 
-`other`
+other
 
 
-A `TypeSpec`.
+A TypeSpec.
 
 

 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If there is no TypeSpec that is compatible with both `self`
-and `other`.
+If there is no TypeSpec that is compatible with both self
+and other.
 
 

 

 
 
relax

 
-View
+View
 source
 
 
@@ -308,30 +296,32 @@ Allows variable number of nodes or/and graph components.
 Calling with all default parameters keeps the spec unchanged.
 
 
+
  

 
 
Args

 
 

 
-`num_components`
+num_components
 
 
 if True, allows variable number of graph components by
-setting the outermost sizes dimension to `None`.
+setting the outermost sizes dimension to None.
 
 

 
-`num_nodes`
+num_nodes
 
 
 if True, allows variable number of nodes by setting the
-outermost features dimensions to `None`.
+outermost features dimensions to None.
 
 

 

 
 
+
  
@@ -344,13 +334,14 @@ Relaxed compatible edge set spec.
 

 
 
Returns

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if edge set is not scalar (rank > 0).
@@ -358,6 +349,45 @@ if edge set is not scalar (rank > 0).
 

 

 
+
with_indices_dtype

+
+View
+source
+
+
+with_indices_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_row_splits_dtype(
+    dtype: tf.dtypes.DType
+) -> 'GraphPieceSpecBase'
+

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

+
+View
+source
+
+
+with_shape(
+    new_shape: ShapeLike
+) -> 'GraphPieceSpecBase'
+

+
+Enforce the common prefix shape on all the contained features.
+
 
__eq__

 
 
@@ -368,10 +398,9 @@ if edge set is not scalar (rank > 0).
 
 Return self==value.
 
-
 
__getitem__

 
-View
+View
 source
 
 
@@ -389,7 +418,3 @@ source
 

 
 Return self!=value.
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/SizeConstraints.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/SizeConstraints.md
index 81329a1a..967e35e3 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/SizeConstraints.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/SizeConstraints.md
@@ -1,17 +1,10 @@
 # tfgnn.SizeConstraints
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Constraints on the number of entities in the graph.
 
@@ -24,45 +17,40 @@ Constraints on the number of entities in the graph.
 )
 

 
-
-
 
-
-
-
-
 
+
  

 
 
Attributes

 
 

 
-`total_num_components`
+total_num_components
 
 
-A `namedtuple` alias for field number 0
+A namedtuple alias for field number 0
 
 

 
-`total_num_nodes`
+total_num_nodes
 
 
-A `namedtuple` alias for field number 1
+A namedtuple alias for field number 1
 
 

 
-`total_num_edges`
+total_num_edges
 
 
-A `namedtuple` alias for field number 2
+A namedtuple alias for field number 2
 
 

 
-`min_nodes_per_component`
+min_nodes_per_component
 
 
-A `namedtuple` alias for field number 3
+A namedtuple alias for field number 3
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/ValidationError.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/ValidationError.md
index 6b6d888f..f7149e21 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/ValidationError.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/ValidationError.md
@@ -1,17 +1,10 @@
 # tfgnn.ValidationError
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A schema validation error.
 
@@ -21,10 +14,7 @@ A schema validation error.
 )
 

 
-
-
 
 
 This exception is raised if in the course of validating the schema for
 correctness some errors are found.
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/add_readout_from_first_node.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/add_readout_from_first_node.md
index 49733d9f..5f745cd5 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/add_readout_from_first_node.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/add_readout_from_first_node.md
@@ -1,17 +1,10 @@
 # tfgnn.add_readout_from_first_node
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Adds a readout structure equivalent to
 tfgnn.gather_first_node().
@@ -35,15 +28,15 @@ Adds a readout structure equivalent to
 
 

 
-`graph`
+graph
 
 
-A scalar `GraphTensor`. If it contains the `readout_node_set`
+A scalar GraphTensor. If it contains the readout_node_set
 already, its size in each graph component must be 1.
 
 

 
-`key`
+key
 
 
 A key, for use with tfgnn.structured_readout(). The input graph
@@ -51,7 +44,7 @@ must not already contain auxiliary edge sets for readout with this key.
 
 

 
-`node_set_name`
+node_set_name
 
 
 The name of the node set from which values are to be read
@@ -59,7 +52,7 @@ out.
 
 

 
-`readout_node_set`
+readout_node_set
 
 
 The name of the auxiliary node set for readout,
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/add_self_loops.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/add_self_loops.md
index 6d865b51..45baa580 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/add_self_loops.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/add_self_loops.md
@@ -1,71 +1,56 @@
 # tfgnn.add_self_loops
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
-Adds self-loops for edge with name `edge_set_name` EVEN if already exist.
+Adds self-loops for `edge_set_name` EVEN if they already exist.
 
 
 tfgnn.add_self_loops(
-    graph: GraphTensor,
-    edge_set_name: gt.EdgeSetName,
-    *,
-    edge_feature_initializer: _EdgeFeatureInitializer = _zero_edge_feat_init
+    graph: GraphTensor, edge_set_name: gt.EdgeSetName
 ) -> GraphTensor
 

 
 
-
-Edge `edge_set_name` must connect pair of nodes of the same node set.
-
 
+
  
-
-
-

 
 
Args

 
 

 
-`graph`
+graph
 
 
-GraphTensor without self-loops. NOTE: If it has self-loops, then
-another round if self-loops will be added.
+A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
-Must connect node pairs of the same node set.
-

-`edge_feature_initializer`
-
-initializes edge features for the self-loop edges.
-It defaults to initializing features of new edges to tf.zeros.
+An edge set in graph that has the same node set as source
+and target.
 
 

 

 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/all_symbols.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/all_symbols.md
index c769ab64..8d458cf3 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/all_symbols.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/all_symbols.md
@@ -11,14 +11,14 @@
 *   tfgnn.ContextSpec
 *   tfgnn.EdgeSet
 *   tfgnn.EdgeSetSpec
-*   tfgnn.Feature
+*   tfgnn.Feature
 *   tfgnn.FeatureDefaultValues
 *   tfgnn.Field
 *   tfgnn.FieldOrFields
 *   tfgnn.FieldSpec
 *   tfgnn.Fields
 *   tfgnn.FieldsSpec
-*   tfgnn.GraphSchema
+*   tfgnn.GraphSchema
 *   tfgnn.GraphTensor
 *   tfgnn.GraphTensorSpec
 *   tfgnn.HyperAdjacency
@@ -47,6 +47,10 @@
 *   tfgnn.create_schema_pb_from_graph_spec
 *   tfgnn.dataset_filter_with_summary
 *   tfgnn.dataset_from_generator
+*   tfgnn.disable_graph_tensor_validation
+*   tfgnn.disable_graph_tensor_validation_at_runtime
+*   tfgnn.enable_graph_tensor_validation
+*   tfgnn.enable_graph_tensor_validation_at_runtime
 *   tfgnn.experimental
 *   tfgnn.find_tight_size_constraints
 *   tfgnn.gather_first_node
@@ -99,7 +103,20 @@
 *   tfgnn.pool
 *   tfgnn.pool_edges_to_context
 *   tfgnn.pool_edges_to_node
+*   tfgnn.pool_neighbors_to_node
+*   tfgnn.pool_neighbors_to_node_feature
 *   tfgnn.pool_nodes_to_context
+*   tfgnn.proto
+*   tfgnn.proto.BigQuery
+*   tfgnn.proto.BigQuery.TableSpec
+*   tfgnn.proto.Context
+*   tfgnn.proto.EdgeSet
+*   tfgnn.proto.Feature
+*   tfgnn.proto.GraphSchema
+*   tfgnn.proto.Metadata
+*   tfgnn.proto.Metadata.KeyValue
+*   tfgnn.proto.NodeSet
+*   tfgnn.proto.OriginInfo
 *   tfgnn.random_graph_tensor
 *   tfgnn.read_schema
 *   tfgnn.reorder_nodes
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_constraints.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_constraints.md
index f2954781..9d002af4 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_constraints.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_constraints.md
@@ -1,17 +1,10 @@
 # tfgnn.assert_constraints
 
-[TOC]
-
 
 
-

 
 
Returns

 

 
-GraphTensor with self-loops added.
+A GraphTensor with self-loops added. A self-loop is added at each node,
+even if some or all of these nodes already have a loop. All feature tensors
+of the edge set are extended to cover the newly added edges with values
+that are all zeros (for numeric features), false (for boolean features), or
+empty (for string features), respectively.
 
 

-
-
+
+ View source
+on GitHub 
 
 Validate the shape constaints of a graph's features at runtime.
 
@@ -21,32 +14,32 @@ Validate the shape constaints of a graph's features at runtime.
 ) -> tf.Operation
 
 
-
-
 
 
 This code returns a TensorFlow op with debugging assertions that ensure the
-parsed data has valid shape constraints for a graph. This can be instantiated
-in your TensorFlow graph while debugging if you believe that your data may be
+parsed data has valid shape constraints for a graph. This can be instantiated in
+your TensorFlow graph while debugging if you believe that your data may be
 incorrectly shaped, or simply applied to a manually produced dataset to ensure
 that those constraints have been applied correctly.
 
 
+
  

 
 
Args

 
 

 
-`graph`
+graph
 
 
-An instance of a `GraphTensor`.
+An instance of a GraphTensor.
 
 

 

 
 
+
  
@@ -57,4 +50,3 @@ A list of check operations.
 

 
 
Returns

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_satisfies_size_constraints.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_satisfies_size_constraints.md
index a2c360b9..eb4bb5bc 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_satisfies_size_constraints.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/assert_satisfies_size_constraints.md
@@ -1,17 +1,10 @@
 # tfgnn.assert_satisfies_size_constraints
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Raises InvalidArgumentError if graph_tensor exceeds size_constraints.
 
@@ -30,8 +23,6 @@ Raises InvalidArgumentError if graph_tensor exceeds size_constraints.
 )
 
 
-
-
 
 
 This function can be used as follows:
@@ -46,26 +37,27 @@ Conceptually, that means this function is like standard tensorflow assertions,
 like `tf.debugging.Assert(satisfies_size_constraints(...))`, but with the
 following important advantages:
 
-- This functions logs a detailed message which size constraint is violated.
-- This function works around a TensorFlow issue to make sure the assertion is
-  executed before the ops it guards, even in the presence of conflicting
-  attempts to eliminate constant subexpressions.
+-   This functions logs a detailed message which size constraint is violated.
+-   This function works around a TensorFlow issue to make sure the assertion is
+    executed before the ops it guards, even in the presence of conflicting
+    attempts to eliminate constant subexpressions.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 a graph tensor to check against target total sizes.
 
 

 
-`size_constraints`
+size_constraints
 
 
 target total sizes for each graph piece.
@@ -74,31 +66,31 @@ target total sizes for each graph piece.
 

 
 
+
  

 
 
Returns

 

 
-Validation operations to execute within a `tf.control_dependencies`.
+Validation operations to execute within a tf.control_dependencies.
 
 

 
 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`tf.errors.InvalidArgumentError`
+tf.errors.InvalidArgumentError
 
 
 if input graph tensor could not be padded to
-the `size_constraints`.
+the size_constraints.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast.md
index a7d40f89..2c7744df 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast.md
@@ -1,17 +1,10 @@
 # tfgnn.broadcast
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Broadcasts values from nodes to edges, or from context to nodes or edges.
 
@@ -47,20 +40,21 @@ of tfgnn.pool(), which comes in
 handy for some algorithms.
 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_homogeneous_node_and_edge_set_name.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_homogeneous_node_and_edge_set_name.md
index 4ef431b1..6dc8de2a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_homogeneous_node_and_edge_set_name.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_homogeneous_node_and_edge_set_name.md
@@ -1,17 +1,10 @@
 # tfgnn.get_homogeneous_node_and_edge_set_name
 
-[TOC]
-
 
 
-

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`from_tag`
+from_tag
 
 
 Values are broadcast from context if this is tfgnn.CONTEXT or
@@ -68,36 +62,36 @@ from the incident node on each edge with this tag.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 The name of the edge set to which values are broadcast, or
-a non-empty sequence of such names. Unless `from_tag=tfgnn.CONTEXT`,
+a non-empty sequence of such names. Unless from_tag=tfgnn.CONTEXT,
 all named edge sets must have the same incident node set at the given tag.
 
 

 
-`node_set_name`
+node_set_name
 
 
 The name of the node set to which values are broadcast,
 or a non-empty sequence of such names. Can only be passed together with
-`from_tag=tfgnn.CONTEXT`. Exactly one of edge_set_name or node_set_name
+from_tag=tfgnn.CONTEXT. Exactly one of edge_set_name or node_set_name
 must be set.
 
 

 
-`feature_value`
+feature_value
 
 
-A tensor of shape `[num_items, *feature_shape]` from which
+A tensor of shape [num_items, *feature_shape] from which
 the broadcast values are taken. The first dimension indexes the items
 from which the broadcast is done (that is, the nodes of the common node
-set identified by `from_tag`, or the graph components in the context).
+set identified by from_tag, or the graph components in the context).
 
 

 
-`feature_name`
+feature_name
 
 
 The name of a feature stored in the graph, for use instead of
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_edges.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_edges.md
index 802c2530..0ee34eee 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_edges.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_edges.md
@@ -1,17 +1,10 @@
 # tfgnn.broadcast_context_to_edges
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Broadcasts a context value to the `edge_set` edges.
 
@@ -38,37 +31,38 @@ shape prefix of the node set being gathered from. One of `feature_value` or
 `feature_name` must be specified.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 An edge set name.
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense graph context feature value. Has a shape
-`[num_components, *feature_shape]`, where `num_components` is the number
-of components in a graph and `feature_shape` is the shape of the feature
+[num_components, *feature_shape], where num_components is the number
+of components in a graph and feature_shape is the shape of the feature
 value for each component.
 
 

 
-`feature_name`
+feature_name
 
 
 A context feature name.
@@ -77,16 +71,16 @@ A context feature name.
 

 
 
+
  

 
 
Returns

 

 
-Graph context value broadcast to the `edge_set` edges. Has a shape
-`[num_edges, *feature_shape]`, where `num_edges` is the number of edges in
-the `edge_set_name` edge set and `feature_shape` is not affected.
+Graph context value broadcast to the edge_set edges. Has a shape
+[num_edges, *feature_shape], where num_edges is the number of edges in
+the edge_set_name edge set and feature_shape is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_nodes.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_nodes.md
index e86e6d53..742a3965 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_nodes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_context_to_nodes.md
@@ -1,17 +1,10 @@
 # tfgnn.broadcast_context_to_nodes
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Broadcasts a context value to the `node_set` nodes.
 
@@ -38,37 +31,38 @@ shape prefix of the node set being gathered from. One of `feature_value` or
 `feature_name` must be specified.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`node_set_name`
+node_set_name
 
 
 A node set name.
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense graph context feature value. Has a shape
-`[num_components, *feature_shape]`, where `num_components` is the number
-of components in a graph and `feature_shape` is the shape of the feature
+[num_components, *feature_shape], where num_components is the number
+of components in a graph and feature_shape is the shape of the feature
 value for each component.
 
 

 
-`feature_name`
+feature_name
 
 
 A context feature name.
@@ -77,16 +71,16 @@ A context feature name.
 

 
 
+
  

 
 
Returns

 

 
-Graph context value broadcast to the `node_set` nodes. Has a shape
-`[num_nodes, *feature_shape]`, where `num_nodes` is the number of nodes in
-the `node_set_name` node set and `feature_shape` is not affected.
+Graph context value broadcast to the node_set nodes. Has a shape
+[num_nodes, *feature_shape], where num_nodes is the number of nodes in
+the node_set_name node set and feature_shape is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_node_to_edges.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_node_to_edges.md
index 5e4f94b9..8b71242d 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_node_to_edges.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/broadcast_node_to_edges.md
@@ -1,17 +1,10 @@
 # tfgnn.broadcast_node_to_edges
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Broadcasts values from nodes to incident edges.
 
@@ -28,40 +21,41 @@ Broadcasts values from nodes to incident edges.
 
 
 
-Given a particular edge set (identified by `edge_set_name` name), this
-operation collects node features from the specific incident node of each edge
-(as indicated by `node_tag`). For example, setting `node_tag=tfgnn.SOURCE` and
+Given a particular edge set (identified by `edge_set_name` name), this operation
+collects node features from the specific incident node of each edge (as
+indicated by `node_tag`). For example, setting `node_tag=tfgnn.SOURCE` and
 `reduce_type='sum'` gathers the source node features over each edge. (See the
 corresponding `pool_edges_to_node()` mirror operation).
 
 The feature to fetch node values from is provided either by name (using
 `feature_name`) and found in the graph tensor itself, or provided explicitly
 (using `feature_value`) in which case its shape has to be compatible with the
-shape prefix of the node set being gathered from. One of `feature_value`
-or `feature_name` must be specified.
+shape prefix of the node set being gathered from. One of `feature_value` or
+`feature_name` must be specified.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 The name of the edge set to which values are broadcast.
 
 

 
-`node_tag`
+node_tag
 
 
 The incident side of each edge from which values are broadcast,
@@ -70,17 +64,17 @@ specified by its tag in the edge set (e.g. tf
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense source node feature values. Has a shape
-`[num_nodes, *feature_shape]`, where `num_nodes` is the number of nodes in
+[num_nodes, *feature_shape], where num_nodes is the number of nodes in
 the incident node set and feature_shape is the shape of the feature value
 for each node.
 
 

 
-`feature_name`
+feature_name
 
 
 A source node feature name.
@@ -89,16 +83,16 @@ A source node feature name.
 

 
 
+
  

 
 
Returns

 

 
-Source node value broadcast to corresponding edges. Has a shape `[num_edges,
-*feature_shape]`, where `num_edges` is the number of edges in the
-`edge_set_name` edge set and `feature_shape` is not affected.
+Source node value broadcast to corresponding edges. Has a shape [num_edges,
+*feature_shape], where num_edges is the number of edges in the
+edge_set_name edge set and feature_shape is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_compatible_with_schema_pb.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_compatible_with_schema_pb.md
index 46759814..c2065d7e 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_compatible_with_schema_pb.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_compatible_with_schema_pb.md
@@ -1,24 +1,17 @@
 # tfgnn.check_compatible_with_schema_pb
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Checks that the given spec or value is compatible with the graph schema.
 
 
 tfgnn.check_compatible_with_schema_pb(
     graph: Union[tfgnn.GraphTensor, tfgnn.GraphTensorSpec],
-    schema: tfgnn.GraphSchema
+    schema: tfgnn.proto.GraphSchema
 ) -> None
 

 
@@ -35,20 +28,21 @@ The `graph` is compatible with the `schema` if
     tfgnn.Adjacency.
 
 
+
  

 
 
Args

 
 

 
-`graph`
+graph
 
 
 The graph tensor or graph tensor spec.
 
 

 
-`schema`
+schema
 
 
 The graph schema.
@@ -57,16 +51,17 @@ The graph schema.
 

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-if `spec_or_value` is not represented by the graph schema.
+if spec_or_value is not represented by the graph schema.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_homogeneous_graph_tensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_homogeneous_graph_tensor.md
index ab00103f..04bfaa9d 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_homogeneous_graph_tensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_homogeneous_graph_tensor.md
@@ -1,17 +1,10 @@
 # tfgnn.check_homogeneous_graph_tensor
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Raises ValueError when tfgnn.get_homogeneous_node_and_edge_set_name() does.
 
@@ -23,4 +16,3 @@ Raises ValueError when tfgnn.get_homogeneous_node_and_edge_set_name() does.
 
 
 
-  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_required_features.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_required_features.md
index b6b546ae..8098e57b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_required_features.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_required_features.md
@@ -1,64 +1,55 @@
 # tfgnn.check_required_features
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Checks the requirements of a given schema against another.
 
 
 tfgnn.check_required_features(
-    requirements: tfgnn.GraphSchema,
-    actual: tfgnn.GraphSchema
+    requirements: tfgnn.proto.GraphSchema,
+    actual: tfgnn.proto.GraphSchema
 )
 

 
-
-
 
 
 This function is used to enable the specification of required features to a
 function. A function accepting a `GraphTensor` instance can this way document
-what features it is expecting to find on it. The function accepts two schemas:
-a `requirements` schema which describes what the function will attempt to
-fetch and use on the `GraphTensor`, and an `actual` schema instance, which is
-the schema describing the dataset. You can use this in your model code to
-ensure that a dataset contains all the expected node sets, edge sets and
-features that the model uses.
+what features it is expecting to find on it. The function accepts two schemas: a
+`requirements` schema which describes what the function will attempt to fetch
+and use on the `GraphTensor`, and an `actual` schema instance, which is the
+schema describing the dataset. You can use this in your model code to ensure
+that a dataset contains all the expected node sets, edge sets and features that
+the model uses.
 
 Note that a dimension with a size of `0` in a feature from the `requirements`
-schema is interpreted specially: it means "accept any value for this
-dimension." The special value `-1` is still used to represent a ragged
-dimension.
+schema is interpreted specially: it means "accept any value for this dimension."
+The special value `-1` is still used to represent a ragged dimension.
 
 (Finally, note that this function predates the existence of `GraphTensorSpec`,
-which is a runtime descriptor for a `GraphTensor`. We may eventually perovide
-an equivalent construct using the `GraphTensorSpec.)
+which is a runtime descriptor for a `GraphTensor`. We may eventually perovide an
+equivalent construct using the `GraphTensorSpec.)
 
 
+
  

 
 
Args

 
 

 
-`requirements`
+requirements
 
 
 An instance of a GraphSchema object, with optional shapes.
 
 

 
-`actual`
+actual
 
 
 The instance of actual schema to check is a matching superset
@@ -68,13 +59,14 @@ of the required schema.
 

 
 
+
  
+    
+    

 
 
Raises

 
 

 
-`ValidationError`
+ValidationError
 
 
 If the given schema does not fulfill the requirements.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_scalar_graph_tensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_scalar_graph_tensor.md
index 59723b3d..fbc17e48 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/check_scalar_graph_tensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_scalar_graph_tensor.md
@@ -1,17 +1,12 @@
 # tfgnn.check_scalar_graph_tensor
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
+
+Checks that graph tensor is scalar (has rank 0).
 
 
 tfgnn.check_scalar_graph_tensor(
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/combine_values.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/combine_values.md
index 6ad1cf5b..ea989c5b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/combine_values.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/combine_values.md
@@ -1,17 +1,10 @@
 # tfgnn.combine_values
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Combines a list of tensors into one (by concatenation or otherwise).
 
@@ -23,32 +16,31 @@ Combines a list of tensors into one (by concatenation or otherwise).
 
 
 
-This is a convenience wrapper around standard TensorFlow operations, to
-provide standard names for common types of combining.
+This is a convenience wrapper around standard TensorFlow operations, to provide
+standard names for common types of combining.
 
 
+
  
 
-
-
-
-

 
 
Args
  `inputs`   a list of Tensors or
+
  inputs   a list of Tensors or
 RaggedTensors, with shapes and types that are compatible for the selected
-combine_type.  
  `combine_type`  one of the following string values, to select the method for combining the
-inputs:
+combine_type.  
 
+combine_type   one of the
+following string values, to select the method for combining the inputs:
 
-  * "sum": The input tensors are added. Their dtypes and shapes must
-    match.
-  * "concat": The input tensors are concatenated along the last axis.
+*   "sum": The input tensors are added. Their dtypes and shapes must match.
+*   "concat": The input tensors are concatenated along the last axis.
     Their dtypes and shapes must match, except for the number of elements
     along the last axis.
-

-
-
+    

 
 
+
  
@@ -59,4 +51,3 @@ A tensor with the combined value of the inputs.
 

 
 
Returns

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/convert_to_line_graph.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/convert_to_line_graph.md
index c53227e6..61e3ab66 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/convert_to_line_graph.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/convert_to_line_graph.md
@@ -1,17 +1,10 @@
 # tfgnn.convert_to_line_graph
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Obtain a graph's line graph.
 
@@ -73,14 +66,14 @@ sets: 'original/to/lines' with edges a->ab, b->bc, c->ca, and
 
 

 
-`graph_tensor`
+graph_tensor
 
 
 Graph to convert to a line graph.
 
 

 
-`connect_from`
+connect_from
 
 
 Specifies which endpoint of the original edges
@@ -88,7 +81,7 @@ will determine the source for the line graph edges.
 
 

 
-`connect_to`
+connect_to
 
 
 Specifies which endpoint of the original edges
@@ -96,29 +89,29 @@ will determine the target for the line graph edges.
 
 

 
-`connect_with_original_nodes`
+connect_with_original_nodes
 
 
 If true, keep the original node sets (not the
 original edge sets) and connect them to line graph nodes according to
 source and target in the original graph. The node set names will be called
-`original/{node_set}` and the new edges `original/to/{edge_set}` for the
-SOURCE nodes and `original/from/{edge_set}` for the TARGET nodes.
+original/{node_set} and the new edges original/to/{edge_set} for the
+SOURCE nodes and original/from/{edge_set} for the TARGET nodes.
 
 

 
-`non_backtracking`
+non_backtracking
 
 
 Whether to return the non-backtracking line graph. Setting
 this to True will only connect edges where the "outer" nodes are
-different, i.e. `u_{1-i} != v_{1-j}`. For default connection settings,
+different, i.e. u_{1-i} != v_{1-j}. For default connection settings,
 for every edge u->v this *removes* line graph edges uv->vu. If
 connect_to=TARGET, this *removes* line graph edges uv->uv.
 
 

 
-`use_node_features_as_line_graph_edge_features`
+use_node_features_as_line_graph_edge_features
 
 
 Whether to use the original
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/create_graph_spec_from_schema_pb.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/create_graph_spec_from_schema_pb.md
index 66b57f8d..71202d63 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/create_graph_spec_from_schema_pb.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/create_graph_spec_from_schema_pb.md
@@ -1,73 +1,64 @@
 # tfgnn.create_graph_spec_from_schema_pb
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Converts a graph schema proto message to a scalar GraphTensorSpec.
 
 
 tfgnn.create_graph_spec_from_schema_pb(
-    schema: tfgnn.GraphSchema,
+    schema: tfgnn.proto.GraphSchema,
     indices_dtype: tf.dtypes.DType = gc.default_indices_dtype
 ) -> tfgnn.GraphTensorSpec
 

 
-
-
 
 
-A `GraphSchema` message contains shape information in a serializable format.
-The `GraphTensorSpec` is a runtime object fulfilling the type spec
-requirements, that accompanies each `GraphTensor` instance and fulfills much
-of the same goal. This function converts the proto to the corresponding type
-spec.
+A `GraphSchema` message contains shape information in a serializable format. The
+`GraphTensorSpec` is a runtime object fulfilling the type spec requirements,
+that accompanies each `GraphTensor` instance and fulfills much of the same goal.
+This function converts the proto to the corresponding type spec.
 
-It is guranteed that the output graph spec is compatible with the input graph
+It is guaranteed that the output graph spec is compatible with the input graph
 schema (as
 tfgnn.check_compatible_with_schema_pb().)
 
 
+
  

 
 
Args

 
 

 
-`schema`
+schema
 
 
 An instance of the graph schema proto message.
 
 

 
-`indices_dtype`
+indices_dtype
 
 
-A `tf.dtypes.DType` for GraphTensor edge set source and
+A tf.dtypes.DType for GraphTensor edge set source and
 target indices, node and edge sets sizes.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A `GraphTensorSpec` specification for the scalar graph tensor (of rank 0).
+A GraphTensorSpec specification for the scalar graph tensor (of rank 0).
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/create_schema_pb_from_graph_spec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/create_schema_pb_from_graph_spec.md
index ac837446..5cd788df 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/create_schema_pb_from_graph_spec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/create_schema_pb_from_graph_spec.md
@@ -1,24 +1,17 @@
 # tfgnn.create_schema_pb_from_graph_spec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Converts scalar GraphTensorSpec to a graph schema proto message.
 
 
 tfgnn.create_schema_pb_from_graph_spec(
     graph: Union[tfgnn.GraphTensor, tfgnn.GraphTensorSpec]
-) -> tfgnn.GraphSchema
+) -> tfgnn.proto.GraphSchema
 

 
 
@@ -29,18 +22,19 @@ features. For edge sets their `source` and `target` fields are populated. All
 other fields are left unset. (Callers can set them separately before writing out
 the schema.)
 
-It is guranteed that the input graph is compatible with the output graph schema
+It is guaranteed that the input graph is compatible with the output graph schema
 (as
 tfgnn.check_compatible_with_schema_pb().)
 
 
+
  

 
 
Args

 
 

 
-`graph`
+graph
 
 
 The scalar graph tensor or its spec with single graph component.
@@ -49,6 +43,7 @@ The scalar graph tensor or its spec with single graph component.
 

 
 
+
  
@@ -61,23 +56,32 @@ An instance of the graph schema proto message.
 

 
 
Returns

 
 
+
  
+
+
+

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if graph has multiple graph components or rank > 0.
 
 

 
-`ValueError`
+ValueError
+
+if adjacency types is not an instance of fgnn.Adjacency.
+

+ValueError
 
 
-if adjacency types is not an instance of `fgnn.Adjacency`.
+if graph tensor features have types that are not supported
+by the graph schema.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_filter_with_summary.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_filter_with_summary.md
index c1fa8af2..4fecb22b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_filter_with_summary.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_filter_with_summary.md
@@ -1,17 +1,10 @@
 # tfgnn.dataset_filter_with_summary
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Dataset filter with a summary for the fraction of dataset elements removed.
 
@@ -26,8 +19,6 @@ Dataset filter with a summary for the fraction of dataset elements removed.
 )
 
 
-
-
 
 
 The fraction of removed elements is computed using exponential moving average.
@@ -36,38 +27,39 @@ See https://en.wikipedia.org/wiki/Moving_average.
 The summary is reported each `summary_steps` elements in the input dataset
 before filtering. Statistics are reported using `tf.summary.scalar()` with
 `step` set to the element index in the result (filtered) dataset, see
-https://tensorflow.org/tensorboard/scalars_and_keras#logging_custom_scalars
-for how to write and retreive them.
+https://tensorflow.org/tensorboard/scalars_and_keras#logging_custom_scalars for
+how to write and retrieve them.
 
 
+
  

 
 
Args

 
 

 
-`dataset`
+dataset
 
 
 An input dataset.
 
 

 
-`predicate`
+predicate
 
 
 A function mapping a dataset element to a boolean.
 
 

 
-`summary_name`
+summary_name
 
 
 A name for this summary.
 
 

 
-`summary_steps`
+summary_steps
 
 
 Report summary for this number of elements in the input
@@ -75,25 +67,25 @@ dataset before filtering.
 
 

 
-`summary_decay`
+summary_decay
 
 
 An exponential moving average decay factor. If not set,
-defaults to the `exp(- 1 / summary_steps)`.
+defaults to the exp(- 1 / summary_steps).
 
 

 

 
 
+
  

 
 
Returns

 

 
 Thed dataset containing the elements of this dataset for which predicate is
-`True`.
+True.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_from_generator.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_from_generator.md
index fa38d480..fb45fa7e 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_from_generator.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/dataset_from_generator.md
@@ -1,17 +1,10 @@
 # tfgnn.dataset_from_generator
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Creates dataset from generator of any nest of scalar graph pieces.
 
@@ -42,13 +35,14 @@ print([dataset1])  # prints: pieceA, pieceD.
 ```
 
 
+
  

 
 
Args

 
 

 
-`generator`
+generator
 
 
 a callable object that returns an object that supports the iter()
@@ -60,25 +54,27 @@ protocol. Could consist of any nest of tensors and scalar graph pieces
 

 
 
+
  

 
 
Returns

 

 
-A `tf.data.Dataset`.
+A tf.data.Dataset.
 
 

 
 

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if any contained graph piece is not scalar or has not compatible
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation.md
new file mode 100644
index 00000000..5532d788
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation.md
@@ -0,0 +1,19 @@
+# tfgnn.disable_graph_tensor_validation
+
+
+
+
+ View source
+on GitHub 
+
+Disables both static and runtime checks of graph tensors.
+
+
+tfgnn.disable_graph_tensor_validation()
+

+
+
+
+IMPORTANT: This is temporary workaround for the legacy code (before TF-GNN 1.0
+release) that may rely on the inconsistent number of graph tensor items and
+allowed edges with adjaceny indices for non-existing nodes. **DO NOT USE**.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation_at_runtime.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation_at_runtime.md
new file mode 100644
index 00000000..160333b7
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/disable_graph_tensor_validation_at_runtime.md
@@ -0,0 +1,15 @@
+# tfgnn.disable_graph_tensor_validation_at_runtime
+
+
+
+
+ View source
+on GitHub 
+
+Disables runtime checks (`tf.debugging.Assert`) of graph tensors.
+
+
+tfgnn.disable_graph_tensor_validation_at_runtime()
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation.md
new file mode 100644
index 00000000..33441682
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation.md
@@ -0,0 +1,15 @@
+# tfgnn.enable_graph_tensor_validation
+
+
+
+
+ View source
+on GitHub 
+
+Enables static checks of graph tensors.
+
+
+tfgnn.enable_graph_tensor_validation()
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation_at_runtime.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation_at_runtime.md
new file mode 100644
index 00000000..7ca92008
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/enable_graph_tensor_validation_at_runtime.md
@@ -0,0 +1,15 @@
+# tfgnn.enable_graph_tensor_validation_at_runtime
+
+
+
+
+ View source
+on GitHub 
+
+Enables both static and runtime checks of graph tensors.
+
+
+tfgnn.enable_graph_tensor_validation_at_runtime()
+

+
+
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md
index 807da7f4..95d30833 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md
@@ -1,17 +1,10 @@
 # Module: tfgnn.experimental
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Experimental (unstable) parts of the public interface of TensorFlow GNN.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/find_tight_size_constraints.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/find_tight_size_constraints.md
index 393fcb07..f81f5624 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/find_tight_size_constraints.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/find_tight_size_constraints.md
@@ -1,17 +1,10 @@
 # tfgnn.find_tight_size_constraints
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns smallest possible size constraints that allow dataset padding.
 
@@ -24,17 +17,15 @@ Returns smallest possible size constraints that allow dataset padding.
 ) -> tfgnn.SizeConstraints
 
 
-
-
 
 
 Evaluated constraints are intended to be used when it is required that all
 elements of the `dataset` can be padded, e.g., when evaluating models.
 
 Typically, this function is used on a dataset of individual examples (that is,
-not batched), and the `target_batch_size` is passed as an argument. The
-returned constraints will work for all possible batches up to that size drawn
-from the dataset.
+not batched), and the `target_batch_size` is passed as an argument. The returned
+constraints will work for all possible batches up to that size drawn from the
+dataset.
 
 Alternatively, this function can be used on a dataset that is already batched,
 passing `target_batch_size=None`. The returned constraints will work for the
@@ -43,24 +34,25 @@ optimized ways of building a Dataset (like parallel .map() and .interleave()
 calls before .batch()) introduce nondeterminism and may not deliver the exact
 same batches again.
 
-Note that this function iterates over all elements of the input dataset, so
-its execution time is proportional to the dataset's cardinality.
+Note that this function iterates over all elements of the input dataset, so its
+execution time is proportional to the dataset's cardinality.
 
 
+
  

 
 
Args

 
 

 
-`dataset`
+dataset
 
 
 finite dataset of graph tensors of any rank.
 
 

 
-`min_nodes_per_component`
+min_nodes_per_component
 
 
 mapping from a node set name to a minimum number of
@@ -68,16 +60,17 @@ nodes in each graph component. Defaults to 0.
 
 

 
-`target_batch_size`
+target_batch_size
 
 
-if not `None`, an integer for multiplying the sizes
+if not None, an integer for multiplying the sizes
 measured from dataset before making room for padding.
 
 

 

 
 
+
  
@@ -90,20 +83,19 @@ in the input dataset.
 
 

 
 
Returns

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if dataset elements are not GraphTensors or its cardinality
-is `tf.data.INFINITE_CARDINALITY`.
+is tf.data.INFINITE_CARDINALITY.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/gather_first_node.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/gather_first_node.md
index 2f380ccc..4067d3f6 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/gather_first_node.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/gather_first_node.md
@@ -1,17 +1,10 @@
 # tfgnn.gather_first_node
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Gathers feature value from the first node of each graph component.
 
@@ -27,54 +20,55 @@ Gathers feature value from the first node of each graph component.
 
 
 
-Given a particular node set (identified by `node_set_name`), this operation
-will gather the given feature from the first node of each graph component.
+Given a particular node set (identified by `node_set_name`), this operation will
+gather the given feature from the first node of each graph component.
 
 This is often used for rooted graphs created by sampling around the
 neighborhoods of seed nodes in a large graph: by convention, each seed node is
 the first node of its component in the respective node set, and this operation
 reads out the information it has accumulated there. (In other node sets, the
-first node may be arbitrary -- or nonexistant, in which case this operation
-must not be used and may raise an error at runtime.)
+first node may be arbitrary -- or nonexistant, in which case this operation must
+not be used and may raise an error at runtime.)
 
 The feature to fetch node values from is provided either by name (using
 `feature_name`) and found in the graph tensor itself, or provided explicitly
 (using `feature_value`) in which case its shape has to be compatible with the
-shape prefix of the node set being gathered from. One of `feature_value`
-or `feature_name` must be specified.
+shape prefix of the node set being gathered from. One of `feature_value` or
+`feature_name` must be specified.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`node_set_name`
+node_set_name
 
 
 A seed node set name.
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense node feature value. Has a shape
-`[num_nodes, *feature_shape]`, where `num_nodes` is the number of nodes in
-the `node_set_name` node set and `feature_shape` is the shape of the
+[num_nodes, *feature_shape], where num_nodes is the number of nodes in
+the node_set_name node set and feature_shape is the shape of the
 feature value for each node.
 
 

 
-`feature_name`
+feature_name
 
 
 A node feature name.
@@ -83,17 +77,17 @@ A node feature name.
 

 
 
+
  

 
 
Returns

 

 
 A tensor of gathered feature values, one for each graph component, like a
-context feature. Has a shape `[num_components, *feature_shape]`, where
-`num_components` is the number of components in a graph and `feature_shape`
+context feature. Has a shape [num_components, *feature_shape], where
+num_components is the number of components in a graph and feature_shape
 is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_aux_type_prefix.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_aux_type_prefix.md
index f5175ed4..c61e57b1 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_aux_type_prefix.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_aux_type_prefix.md
@@ -1,17 +1,10 @@
 # tfgnn.get_aux_type_prefix
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns type prefix of aux node or edge set names, or `None` if non-aux.
 
@@ -52,11 +45,11 @@ and node sets are formed.
 
 

 
-`set_name`
+set_name
 
 
 The name of a node set or edge set in a tfgnn.GraphTensor,
-tfgnn.GraphTensorSpec or tfgnn.GraphSchema.
+tfgnn.GraphTensorSpec or tfgnn.GraphSchema.
 
 

 

@@ -69,7 +62,7 @@ The name of a node set or edge set in a 
 

 
 For an auxiliary node set or edge set, a non-empty prefix that identifies
-its type; for other node sets or edge sets, `None`.
+its type; for other node sets or edge sets, None.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns the sole `node_set_name, edge_set_name` or raises `ValueError`.
 
@@ -38,14 +31,14 @@ as appropriate for model-building code.
 
 

 
-`graph`
+graph
 
 
-the `GraphTensor` or `GraphTensorSpec` to check.
+the GraphTensor or GraphTensorSpec to check.
 
 

 
-`name`
+name
 
 
 Optionally, the name of the operation (library function, class, ...)
@@ -62,9 +55,9 @@ raised.
 
Returns

 

 
-A tuple `node_set_name, edge_set_name` with the unique node set and edge
-set, resp., in `graph` for which tfgnn.get_aux_type_prefix(set_name) is
-`None`.
+A tuple node_set_name, edge_set_name with the unique node set and edge
+set, resp., in graph for which tfgnn.get_aux_type_prefix(set_name) is
+None.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns tf.io parsing features for `GraphTensorSpec` type spec.
 
@@ -27,30 +20,32 @@ Returns tf.io parsing features for `GraphTensorSpec` type spec.
 
 This function returns a mapping of `tf.train.Feature` names to configuration
 objects that can be used to parse instances of `tf.train.Example` (see
-https://www.tensorflow.org/api_docs/python/tf/io). The resulting mapping can
-be used with `tf.io.parse_example()` for reading the individual fields of a
+https://www.tensorflow.org/api_docs/python/tf/io). The resulting mapping can be
+used with `tf.io.parse_example()` for reading the individual fields of a
 `GraphTensor` instance. This essentially forms our encoding of a `GraphTensor`
 to a `tf.train.Example` proto.
 
 (This is an internal function. You are not likely to be using this if you're
-decoding graph tensors. Instead, you should use the tfgnn.parse_example()
+decoding graph tensors. Instead, you should use the
+tfgnn.parse_example()
 routine directly, which handles this process for you.)
 
 
+
  

 
 
Args

 
 

 
-`spec`
+spec
 
 
 A graph tensor type specification.
 
 

 
-`prefix`
+prefix
 
 
 An optional prefix string over all the features. You may use
@@ -58,25 +53,25 @@ this if you are encoding other data in the same protocol buffer.
 
 

 
-`validate`
+validate
 
 
 A boolean indicating whether or not to validate that the input
-fields form a valid `GraphTensor`. Defaults to `True`.
+fields form a valid GraphTensor. Defaults to True.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A dict of `tf.train.Feature` name to feature configuration object, to be
-used in `tf.io.parse_example()`.
+A dict of tf.train.Feature name to feature configuration object, to be
+used in tf.io.parse_example().
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_registered_reduce_operation_names.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_registered_reduce_operation_names.md
index 12421892..e283e116 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/get_registered_reduce_operation_names.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/get_registered_reduce_operation_names.md
@@ -1,17 +1,10 @@
 # tfgnn.get_registered_reduce_operation_names
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns the registered list of supported reduce operation names.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/graph_tensor_to_values.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/graph_tensor_to_values.md
index e5defd8e..0cab465a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/graph_tensor_to_values.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/graph_tensor_to_values.md
@@ -1,17 +1,10 @@
 # tfgnn.graph_tensor_to_values
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Convert an eager `GraphTensor` to a mapping of mappings of PODTs.
 
@@ -21,8 +14,6 @@ Convert an eager `GraphTensor` to a mapping of mappings of PODTs.
 ) -> Dict[str, Any]
 
 
-
-
 
 
 This is used for pretty-printing. Convert your graph tensor with this and run
@@ -30,30 +21,31 @@ the result through `pprint.pprint()` or `pprint.pformat()` for display of its
 contents.
 
 
+
  

 
 
Args

 
 

 
-`graph`
+graph
 
 
-An eager `GraphTensor` instance to be pprinted.
+An eager GraphTensor instance to be pprinted.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A dict of plain-old data types that can be run through `pprint.pprint()` or
+A dict of plain-old data types that can be run through pprint.pprint() or
 a JSON conversion library.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md
index c74bf357..4a9cc45a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md
@@ -1,17 +1,10 @@
 # tfgnn.homogeneous
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Constructs a homogeneous `GraphTensor` with node features and one edge_set.
 
@@ -31,29 +24,29 @@ Constructs a homogeneous `GraphTensor` with node features and one edge_set.
 
 
 
-
 
+
  

 
 
Args

 
 

 
-`source`
+source
 
 
 A dense Tensor with the source indices for edges
 
 

 
-`target`
+target
 
 
 A dense Tensor with the target indices for edges
 
 

 
-`node_features`
+node_features
 
 
 A Tensor or mapping from feature name to Tensor of features
@@ -61,7 +54,7 @@ corresponding to graph nodes.
 
 

 
-`edge_features`
+edge_features
 
 
 Optional Tensor or mapping from feature name to Tensor of
@@ -69,7 +62,7 @@ features corresponding to graph edges.
 
 

 
-`context_features`
+context_features
 
 
 Optional Tensor or mapping from name to Tensor for the
@@ -77,21 +70,21 @@ context (entire graph)
 
 

 
-`node_set_name`
+node_set_name
 
 
 Optional name for the node set
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 Optional name for the edge set
 
 

 
-`node_set_sizes`
+node_set_sizes
 
 
 Optional Tensor with the number of nodes per component. If
@@ -100,7 +93,7 @@ the same length.
 
 

 
-`edge_set_sizes`
+edge_set_sizes
 
 
 Optional Tensor with the number of edges per component. If
@@ -111,12 +104,13 @@ the same length.
 

 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/is_dense_tensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/is_dense_tensor.md
index 82d78186..9150fe64 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/is_dense_tensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/is_dense_tensor.md
@@ -1,17 +1,10 @@
 # tfgnn.is_dense_tensor
 
-[TOC]
-
 
 
-

 
 
Returns

 

 
-A scalar `GraphTensor` with a single node set and edge set.
+A scalar GraphTensor with a single node set and edge set.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns whether a tensor (TF or Keras) is a Tensor.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/is_graph_tensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/is_graph_tensor.md
index 1ae673d4..462cca77 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/is_graph_tensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/is_graph_tensor.md
@@ -1,17 +1,10 @@
 # tfgnn.is_graph_tensor
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns whether `value` is a GraphTensor (possibly wrapped for Keras).
 
@@ -21,6 +14,4 @@ Returns whether `value` is a GraphTensor (possibly wrapped for Keras).
 ) -> bool
 
 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/is_ragged_tensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/is_ragged_tensor.md
index d534e1e6..258bd27f 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/is_ragged_tensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/is_ragged_tensor.md
@@ -1,17 +1,10 @@
 # tfgnn.is_ragged_tensor
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns whether a tensor (TF or Keras) is a RaggedTensor.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_features.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_features.md
index 2e582687..ed1f03eb 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_features.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_features.md
@@ -1,49 +1,42 @@
 # tfgnn.iter_features
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Utility function to iterate over the features of a graph schema.
 
 
 tfgnn.iter_features(
-    schema: Union[tfgnn.GraphSchema, tfgnn.GraphTensor]
+    schema: Union[tfgnn.proto.GraphSchema, tfgnn.GraphTensor]
 ) -> Iterator[Tuple[Text, Text, Text, Union[schema_pb2.Feature, gt.Field]]]
 

 
-
-
 
 
 This function iterates over all the feature values of each of the context set,
 each of the node sets, and each of the edge sets.
 
 
+
  

 
 
Args

 
 

 
-`schema`
+schema
 
 
-An instance of a `GraphSchema` proto message.
+An instance of a GraphSchema proto message.
 
 

 

 
 
+
  
@@ -51,13 +44,10 @@ An instance of a `GraphSchema` proto message.
 
-
+*   set-type: A type of set, which is either of "context", "nodes" or "edges".
+*   set-name: A string, the name of the set.
+*   feature-name: A string, the name of the feature in the set.
+*   feature-value: A potentially ragged tensor (either a tf.Tensor
+    or a tf.RaggedTensor). 

 
 
Yields

 Triplets of (set-type, set-name, feature-name, feature-value) where
 
-* set-type: A type of set, which is either of "context", "nodes" or "edges".
-* set-name: A string, the name of the set.
-* feature-name: A string, the name of the feature in the set.
-* feature-value: A potentially ragged tensor (either a `tf.Tensor` or a
-  `tf.RaggedTensor`).
-
 
 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_sets.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_sets.md
index a0ca5a66..8941eb7a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_sets.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/iter_sets.md
@@ -1,49 +1,42 @@
 # tfgnn.iter_sets
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Utility function to iterate over all the sets present in a graph schema.
 
 
 tfgnn.iter_sets(
-    schema: Union[tfgnn.GraphSchema, tfgnn.GraphTensor]
+    schema: Union[tfgnn.proto.GraphSchema, tfgnn.GraphTensor]
 ) -> Iterator[Tuple[str, str, Any]]
 

 
-
-
 
 
-This function iterates over the context set, each of the node sets, and
-finally each of the edge sets.
+This function iterates over the context set, each of the node sets, and finally
+each of the edge sets.
 
 
+
  

 
 
Args

 
 

 
-`schema`
+schema
 
 
-An instance of a `GraphSchema` proto message.
+An instance of a GraphSchema proto message.
 
 

 

 
 
+
  
@@ -51,11 +44,8 @@ An instance of a `GraphSchema` proto message.
 
-
+*   set-type: A type of set, which is either of "context", "nodes" or "edges".
+*   set-name: A string, the name of the set.
+*   features: A dict of feature-name to feature-value. 

 
 
Yields

 Triplets of (set-type, set-name, features) where
 
-* set-type: A type of set, which is either of "context", "nodes" or "edges".
-* set-name: A string, the name of the set.
-* features: A dict of feature-name to feature-value.
-
 
 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras.md
index 257ea6c4..99041b1b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras.md
@@ -1,31 +1,21 @@
 # Module: tfgnn.keras
 
-[TOC]
-
 
 
-
-
-
-
-
+
+ View source
+on GitHub 
 
 The tfgnn.keras package.
 
-
-
 ## Modules
 
 [`layers`](../tfgnn/keras/layers.md) module: The tfgnn.keras.layers package.
 
 ## Classes
 
-[`class ConvGNNBuilder`](../tfgnn/keras/ConvGNNBuilder.md): Factory of layers that do convolutions on a graph.
+[`class ConvGNNBuilder`](../tfgnn/keras/ConvGNNBuilder.md): Factory of layers
+that do convolutions on a graph.
 
 ## Functions
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/ConvGNNBuilder.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/ConvGNNBuilder.md
index 368fcf70..f99101fe 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/ConvGNNBuilder.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/ConvGNNBuilder.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.ConvGNNBuilder
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Factory of layers that do convolutions on a graph.
 
@@ -31,9 +24,8 @@ Factory of layers that do convolutions on a graph.
 ConvGNNBuilder object constructs `GraphUpdate` layers, that apply arbitrary
 convolutions and updates on nodes of a graph. The convolutions (created by the
 `convolutions_factory`) propagate information to the incident edges of the
-graph. The results of the convolution together with the current nodes states
-are used to update the nodes, using a layer created by
-`nodes_next_state_factory`.
+graph. The results of the convolution together with the current nodes states are
+used to update the nodes, using a layer created by `nodes_next_state_factory`.
 
 Layers created by ConvGNNBuilder can be (re-)used in any order.
 
@@ -75,33 +67,34 @@ def graph_update_factory(deferred_init_callback, name):
 ```
 
 
+
  
@@ -147,20 +140,21 @@ source
 
 Constructs GraphUpdate layer for the set of receiver node sets.
 
-This method contructs NodeSetUpdate layers from convolutions and next state
+This method constructs NodeSetUpdate layers from convolutions and next state
 factories (specified during the class construction) for the given receiver node
 sets. The resulting node set update layers are combined and returned as one
 GraphUpdate layer. Auxiliary node sets (e.g., as needed for
 `tfgnn.keras.layers.NamedReadout`) are ignored.
 
 
+
  

 
 
Init args

 
 

 
-`convolutions_factory`
+convolutions_factory
 
 
 called as
-`convolutions_factory(edge_set_name, receiver_tag=receiver_tag)`
+convolutions_factory(edge_set_name, receiver_tag=receiver_tag)
 to return the convolution layer for the edge set towards the specified
-receiver. The `receiver_tag` kwarg is omitted from the call if it is
+receiver. The receiver_tag kwarg is omitted from the call if it is
 omitted from the init args (but that usage is deprecated).
 
 

 
-`nodes_next_state_factory`
+nodes_next_state_factory
 
 
 called as
-`nodes_next_state_factory(node_set_name)` to return the next-state layer
+nodes_next_state_factory(node_set_name) to return the next-state layer
 for the respectve NodeSetUpdate.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
 Set this to tfgnn.TARGET or tfgnn.SOURCE to choose which
@@ -111,21 +104,21 @@ New code is expected to set it in any case.
 
 

 
-`node_set_update_factory`
+node_set_update_factory
 
 
 If set, called as
-`node_set_update_factory(node_set_name, edge_set_inputs, next_state)`
-to return the node set update for the given `node_set_name`. The
+node_set_update_factory(node_set_name, edge_set_inputs, next_state)
+to return the node set update for the given node_set_name. The
 remaining arguments are as expected by tfgnn.keras.layers.NodeSetUpdate.
 
 

 
-`graph_update_factory`
+graph_update_factory
 
 
 If set, called as
-`graph_update_factory(deferred_init_callback, name)` to return the graph
+graph_update_factory(deferred_init_callback, name) to return the graph
 update. The arguments are as expected by tfgnn.keras.layers.GraphUpdate.
 
 

 
 
Args

 
 

 
-`node_sets`
+node_sets
 
 
 By default, the result updates all node sets that receive from
@@ -171,15 +165,16 @@ auxiliary node sets.
 
 

 
-`name`
+name
 
 
-Optionally, a name for the returned GraphUpate layer.
+Optionally, a name for the returned GraphUpdate layer.
 
 

 

 
 
+
  
@@ -190,8 +185,3 @@ A GraphUpdate layer, with building deferred to the first call.
 

 
 
Returns

 
 

-
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/clone_initializer.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/clone_initializer.md
index f8deffef..22ecb47a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/clone_initializer.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/clone_initializer.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.clone_initializer
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Clones an initializer to ensure a new default seed.
 
@@ -32,9 +25,9 @@ seed (even if not explicitly specified) at creation time, so that all calls to
 them return the same sequence of numbers. To achieve independent initializations
 of the various model weights, user-specified initializers must be cloned for
 each weight before passing them to Keras. This way, each of them gets a separate
-seed (unless explicitly overriden).
+seed (unless explicitly overridden).
 
-This helper function clones `Initializer` obejcts and passes through all other
+This helper function clones `Initializer` objects and passes through all other
 forms of specifying an initializer. TF-GNN's modeling code applies it before
 passing user-specified initaializers to Keras. User code that calls Keras
 directly and passes an initializer more than once is advised to wrap it with
@@ -66,7 +59,7 @@ return gnn_builder.Convolve()
 
 

 
-`initializer`
+initializer
 
 
 An initializer specification as understood by Keras.
@@ -81,8 +74,8 @@ An initializer specification as understood by Keras.
 
Returns

 

 
-A new `Initializer` object with the same config as `initializer`,
-or `initializer` unchanged if if was not an `Initializer` object.
+A new Initializer object with the same config as initializer,
+or initializer unchanged if it was not an Initializer object.
 
 

-
-
-
-
+
+ View source
+on GitHub 
 
 The tfgnn.keras.layers package.
 
-
-
 ## Classes
 
 [`class AddReadoutFromFirstNode`](../../tfgnn/keras/layers/AddReadoutFromFirstNode.md):
@@ -28,36 +17,49 @@ Adds readout node set equivalent to
 [`class AddSelfLoops`](../../tfgnn/keras/layers/AddSelfLoops.md): Adds
 self-loops to scalar graphs.
 
-[`class AnyToAnyConvolutionBase`](../../tfgnn/keras/layers/AnyToAnyConvolutionBase.md): Convenience base class for convolutions to nodes or to context.
+[`class AnyToAnyConvolutionBase`](../../tfgnn/keras/layers/AnyToAnyConvolutionBase.md):
+Convenience base class for convolutions to nodes or to context.
 
-[`class Broadcast`](../../tfgnn/keras/layers/Broadcast.md): Broadcasts a GraphTensor feature.
+[`class Broadcast`](../../tfgnn/keras/layers/Broadcast.md): Broadcasts a
+GraphTensor feature.
 
-[`class ContextUpdate`](../../tfgnn/keras/layers/ContextUpdate.md): A context update with input from node sets and/or edge sets.
+[`class ContextUpdate`](../../tfgnn/keras/layers/ContextUpdate.md): A context
+update with input from node sets and/or edge sets.
 
-[`class EdgeSetUpdate`](../../tfgnn/keras/layers/EdgeSetUpdate.md): Computes the new state of an EdgeSet from select input features.
+[`class EdgeSetUpdate`](../../tfgnn/keras/layers/EdgeSetUpdate.md): Computes the
+new state of an EdgeSet from select input features.
 
-[`class GraphUpdate`](../../tfgnn/keras/layers/GraphUpdate.md): Applies one round of updates to EdgeSets, NodeSets and Context.
+[`class GraphUpdate`](../../tfgnn/keras/layers/GraphUpdate.md): Applies one
+round of updates to EdgeSets, NodeSets and Context.
 
 [`class ItemDropout`](../../tfgnn/keras/layers/ItemDropout.md): Dropout of
 feature values for entire edges, nodes or components.
 
-[`class MakeEmptyFeature`](../../tfgnn/keras/layers/MakeEmptyFeature.md): Returns an empty feature with a shape that fits the input graph piece.
+[`class MakeEmptyFeature`](../../tfgnn/keras/layers/MakeEmptyFeature.md):
+Returns an empty feature with a shape that fits the input graph piece.
 
-[`class MapFeatures`](../../tfgnn/keras/layers/MapFeatures.md): Transforms features on a GraphTensor by user-defined callbacks.
+[`class MapFeatures`](../../tfgnn/keras/layers/MapFeatures.md): Transforms
+features on a GraphTensor by user-defined callbacks.
 
-[`class NextStateFromConcat`](../../tfgnn/keras/layers/NextStateFromConcat.md): Computes a new state by concatenating inputs and applying a Keras Layer.
+[`class NextStateFromConcat`](../../tfgnn/keras/layers/NextStateFromConcat.md):
+Computes a new state by concatenating inputs and applying a Keras Layer.
 
-[`class NodeSetUpdate`](../../tfgnn/keras/layers/NodeSetUpdate.md): A node state update with input from convolutions or other edge set inputs.
+[`class NodeSetUpdate`](../../tfgnn/keras/layers/NodeSetUpdate.md): A node state
+update with input from convolutions or other edge set inputs.
 
-[`class PadToTotalSizes`](../../tfgnn/keras/layers/PadToTotalSizes.md): Applies tfgnn.pad_to_total_sizes() to a GraphTensor.
+[`class PadToTotalSizes`](../../tfgnn/keras/layers/PadToTotalSizes.md): Applies
+tfgnn.pad_to_total_sizes() to a GraphTensor.
 
-[`class ParseExample`](../../tfgnn/keras/layers/ParseExample.md): Applies tfgnn.parse_example(graph_tensor_spec, _) to a batch of strings.
+[`class ParseExample`](../../tfgnn/keras/layers/ParseExample.md): Applies
+tfgnn.parse_example(graph_tensor_spec, _) to a batch of strings.
 
-[`class ParseSingleExample`](../../tfgnn/keras/layers/ParseSingleExample.md): Applies tfgnn.parse_single_example(graph_tensor_spec, _).
+[`class ParseSingleExample`](../../tfgnn/keras/layers/ParseSingleExample.md):
+Applies tfgnn.parse_single_example(graph_tensor_spec, _).
 
 [`class Pool`](../../tfgnn/keras/layers/Pool.md): Pools a GraphTensor feature.
 
-[`class Readout`](../../tfgnn/keras/layers/Readout.md): Reads a feature out of a GraphTensor.
+[`class Readout`](../../tfgnn/keras/layers/Readout.md): Reads a feature out of a
+GraphTensor.
 
 [`class ReadoutFirstNode`](../../tfgnn/keras/layers/ReadoutFirstNode.md): Reads
 a feature from the first node of each graph component.
@@ -65,7 +67,8 @@ a feature from the first node of each graph component.
 [`class ReadoutNamedIntoFeature`](../../tfgnn/keras/layers/ReadoutNamedIntoFeature.md):
 Reads out a feature value from select nodes (or edges) in a graph.
 
-[`class ResidualNextState`](../../tfgnn/keras/layers/ResidualNextState.md): Updates a state with a residual block.
+[`class ResidualNextState`](../../tfgnn/keras/layers/ResidualNextState.md):
+Updates a state with a residual block.
 
 [`class SimpleConv`](../../tfgnn/keras/layers/SimpleConv.md): A convolution
 layer that applies a passed-in message_fn.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddReadoutFromFirstNode.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddReadoutFromFirstNode.md
index a5ce1ee2..99f6c175 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddReadoutFromFirstNode.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AddReadoutFromFirstNode.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.AddReadoutFromFirstNode
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Adds readout node set equivalent to
 tfgnn.keras.layers.ReadoutFirstNode.
@@ -35,7 +28,7 @@ Adds readout node set equivalent to
 
 

 
-`key`
+key
 
 
 A key, for use with tfgnn.keras.layers.StructuredReadout. The input
@@ -44,7 +37,7 @@ key.
 
 

 
-`node_set_name`
+node_set_name
 
 
 The name of the node set from which values are to be read
@@ -52,7 +45,7 @@ out.
 
 

 
-`readout_node_set`
+readout_node_set
 
 
 The name of the auxiliary node set for readout,
@@ -69,10 +62,10 @@ as in tfgnn.structured_read
 
 

 
-`graph`
+graph
 
 
-A scalar `GraphTensor`. If it contains the readout_node_set already,
+A scalar GraphTensor. If it contains the readout_node_set already,
 its size in each graph component must be 1.
 
 
Returns

 

 
-A modified `GraphTensor` so that tfgnn.keras.layers.StructuredReadout(key)
+A modified GraphTensor so that tfgnn.keras.layers.StructuredReadout(key)
 acts like tfgnn.keras.layers.ReadoutFirstNode(node_set_name=node_set_name)
 on the input graph.
 

-
-
+
+ View source
+on GitHub 
 
 Adds self-loops to scalar graphs.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AnyToAnyConvolutionBase.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AnyToAnyConvolutionBase.md
index dfc89475..295b5ca3 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AnyToAnyConvolutionBase.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/AnyToAnyConvolutionBase.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.AnyToAnyConvolutionBase
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Convenience base class for convolutions to nodes or to context.
 
@@ -27,8 +20,6 @@ Convenience base class for convolutions to nodes or to context.
 )
 

 
-
-
 
 
 This base class simplifies the implementation of graph convolutions as Keras
@@ -36,10 +27,10 @@ Layers. Instead of subclassing Layer directly and implementing `call()` with a
 GraphTensor input, users can subclass this class and implement the abstract
 `convolve()` method, which is invoked with the relevant tensors unpacked from
 the graph, and with callbacks to broadcast and pool values across the relevant
-parts of the graph structure (see its docstring for more).  The resulting
+parts of the graph structure (see its docstring for more). The resulting
 subclass can perform convolutions from nodes to nodes, optionally with a side
-input from edges, or the equivalent computations from nodes to context,
-from edges to context, or from edges into incident nodes.
+input from edges, or the equivalent computations from nodes to context, from
+edges to context, or from edges into incident nodes.
 
 Here is a minimal example:
 
@@ -72,63 +63,51 @@ class ExampleConvolution(tfgnn.keras.layers.AnyToAnyConvolutionBase):
 
 The resulting subclass can be applied to a GraphTensor in four ways:
 
- 1. Convolutions from nodes.
-
-    a) The classic case: convolutions over an edge set.
-       ```
-         sender_node -> sender_edge <->> receiver (node)
-       ```
-       A message is computed for each edge of the edge set, depending on
-       inputs from the sender node, the receiver node (single arrowheads),
-       and/or the edge itself. The messages are aggregated at each receiver
-       node from its incident edges (double arrowhead).
-
-    b) Convolutions from a node set to context.
-       ```
-         sender_node <->> receiver (context)
-       ```
-       Instead of the EdgeSet as in case (1a), there is a NodeSet, which
-       tracks the containment of its nodes in graph components. Instead of
-       a receiver node, there is the per-component graph context.
-       A message is computed for each node of the node set, depending on
-       inputs from the node itself and possibly its context. The messages are
-       aggregated for each component from the nodes contained in it.
-
-
- 2. Pooling from edges.
-
-    a) Pooling from an edge set to an incident node set.
-       ```
-         sender_edge <->> receiver (node)
-       ```
-       This works like a convolution in case (1a) with the side input for the
-       edge feature switched on and the main input from the sender node
-       switched off.
-
-    b) Pooling from an edge set to context.
-       ```
-         sender_edge <->> receiver (context)
-       ```
-       Like in case (1b), the receiver is the context, and senders are
-       connected to it by containment in a graph component. Unlike case (1b),
-       but similar to case (2a), the sender input comes from an edge set,
-       not a node set.
+1.  Convolutions from nodes.
+
+    a) The classic case: convolutions over an edge set. `sender_node ->
+    sender_edge <->> receiver (node)` A message is computed for each edge of the
+    edge set, depending on inputs from the sender node, the receiver node
+    (single arrowheads), and/or the edge itself. The messages are aggregated at
+    each receiver node from its incident edges (double arrowhead).
+
+    b) Convolutions from a node set to context. `sender_node <->> receiver
+    (context)` Instead of the EdgeSet as in case (1a), there is a NodeSet, which
+    tracks the containment of its nodes in graph components. Instead of a
+    receiver node, there is the per-component graph context. A message is
+    computed for each node of the node set, depending on inputs from the node
+    itself and possibly its context. The messages are aggregated for each
+    component from the nodes contained in it.
+
+2.  Pooling from edges.
+
+    a) Pooling from an edge set to an incident node set. `sender_edge <->>
+    receiver (node)` This works like a convolution in case (1a) with the side
+    input for the edge feature switched on and the main input from the sender
+    node switched off.
+
+    b) Pooling from an edge set to context. `sender_edge <->> receiver
+    (context)` Like in case (1b), the receiver is the context, and senders are
+    connected to it by containment in a graph component. Unlike case (1b), but
+    similar to case (2a), the sender input comes from an edge set, not a node
+    set.
 
 Case (1) is solved directly by subclasses of this class with default init args
-`sender_node_feature=tfgnn.HIDDEN_STATE` and `sender_edge_feature=None`.
-The sub-cases are distinguised by passing `receiver_tag=tfgnn.SOURCE` or
-tfgnn.TARGET for (a) and tfgnn.CONTEXT for (b).
-The side input from edges can be activated in case (1a) by setting the init
-arg `sender_edge_feature=`.
-
-Case (2) is solved indirectly by wrapping a subclass to invert the selection
-for sender inputs: nodes off, edges on. The sub-cases (a) and (b) are again
+`sender_node_feature=tfgnn.HIDDEN_STATE` and `sender_edge_feature=None`. The
+sub-cases are distinguised by passing `receiver_tag=tfgnn.SOURCE` or
+tfgnn.TARGET for (a) and
+tfgnn.CONTEXT for (b). The
+side input from edges can be activated in case (1a) by setting the init arg
+`sender_edge_feature=`.
+
+Case (2) is solved indirectly by wrapping a subclass to invert the selection for
+sender inputs: nodes off, edges on. The sub-cases (a) and (b) are again
 distinguished by the `receiver_tag`. TF-GNN avoids the term "Convolution" for
 this operation and calls it "EdgePool" instead, to emphasize that sender nodes
 are not involved. Whenever it makes sense to repurpose a convolution as an
 EdgePool operation, we recommend that the convolution is accompanied by a
-wrapper function in the following style so that the proper term "edge pool"
-can be used in code:
+wrapper function in the following style so that the proper term "edge pool" can
+be used in code:
 
 ```python
 def ExampleEdgePool(*args, sender_feature=tfgnn.HIDDEN_STATE, **kwargs):
@@ -142,11 +121,12 @@ attention to neighbor nodes, but in this way we can reuse the same code for
 attention to incident edges, or to all nodes/edges in a graph component.
 
 
+
  
 
-
+

 
 
Args
  `receiver_tag`   one of
+
  receiver_tag   one of
 tfgnn.SOURCE,
 tfgnn.TARGET or
 tfgnn.CONTEXT. The results
@@ -158,45 +138,46 @@ of the edges. If set to
 tfgnn.CONTEXT, the layer
 can be called for an edge set or a node set and will aggregate results for
 context (per graph component). If left unset for init, the tag must be passed at
-call time.  
  `receiver_feature`
-  The name of the feature that is read from the receiver graph piece
-and passed as convolve(receiver_input=...).  
 
-`sender_node_feature`   The name of the
-feature that is read from the sender nodes, if any, and passed as
-convolve(sender_node_input=...). NOTICE this must be `None` for use with
-`receiver_tag=tfgnn.CONTEXT` on an edge set, or for pooling from edges without
-sender node states.  
 
-`sender_edge_feature`   The name of the
-feature that is read from the sender edges, if any, and passed as
-convolve(sender_edge_input=...). NOTICE this must not be `None` for use with
-`receiver_tag=tfgnn.CONTEXT` on an edge set.  
 
-`extra_receiver_ops`   A str-keyed
-dictionary of Python callables that are wrapped to bind some arguments and then
-passed on to `convolve()`. Sample usage: `extra_receiver_ops={"softmax":
-tfgnn.softmax}`. The values passed in this dict must be callable as follows,
-with two positional arguments:
+call time.  
 
+receiver_feature   The name
+of the feature that is read from the receiver graph piece and passed as
+convolve(receiver_input=...).  
 
+sender_node_feature   The
+name of the feature that is read from the sender nodes, if any, and passed as
+convolve(sender_node_input=...). NOTICE this must be None for use
+with receiver_tag=tfgnn.CONTEXT on an edge set, or for pooling from
+edges without sender node states.  
 
+sender_edge_feature   The
+name of the feature that is read from the sender edges, if any, and passed as
+convolve(sender_edge_input=...). NOTICE this must not be None for
+use with receiver_tag=tfgnn.CONTEXT on an edge set.  
 extra_receiver_ops   A
+str-keyed dictionary of Python callables that are wrapped to bind some arguments
+and then passed on to convolve(). Sample usage:
+extra_receiver_ops={"softmax": tfgnn.softmax}. The values passed in
+this dict must be callable as follows, with two positional arguments:
 
 ```python
 f(graph, receiver_tag, node_set_name=..., feature_value=..., ...)
 f(graph, receiver_tag, edge_set_name=..., feature_value=..., ...)
 ```
 
-The wrapped callables seen by `convolve()` can be called like
+The wrapped callables seen by convolve() can be called like
 
 ```python
 wrapped_f(feature_value, ...)
 ```
 
-The first three arguments of `f` are set to the input GraphTensor of
+The first three arguments of f are set to the input GraphTensor of
 the layer and the tag/name pair required by tfgnn.broadcast() and
 tfgnn.pool() to move values between the receiver and the messages that
 are computed inside the convolution. The sole positional argument of
-`wrapped_f()` is passed to `f()`  as `feature_value=`, and any keyword
+wrapped_f() is passed to f()  as feature_value=, and any keyword
 arguments are forwarded.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Forwarded to the base class tf.keras.layers.Layer.
@@ -205,30 +186,31 @@ Forwarded to the base class tf.keras.layers.Layer.
 

 
 
+
  

 
 
Attributes

 
 

 
-`takes_receiver_input`
+takes_receiver_input
 
 
-If `False`, all calls to convolve() will get `receiver_input=None`.
+If False, all calls to convolve() will get receiver_input=None.
 
 

 
-`takes_sender_edge_input`
+takes_sender_edge_input
 
 
-If `False`, all calls to convolve() will get `sender_edge_input=None`.
+If False, all calls to convolve() will get sender_edge_input=None.
 
 

 
-`takes_sender_node_input`
+takes_sender_node_input
 
 
-If `False`, all calls to convolve() will get `sender_node_input=None`.
+If False, all calls to convolve() will get sender_node_input=None.
 
 

 

@@ -257,52 +239,52 @@ source
 
 Returns the convolution result.
 
-The Tensor inputs to this function still have their original shapes
-and need to be broadcast such that the leading dimension is indexed
-by the items in the graph for which messages are computed (usually edges;
-except when convolving from nodes to context). In the end, values have to be
-pooled from there into a Tensor with a leading dimension indexed by
-receivers, see `pool_to_receiver`.
+The Tensor inputs to this function still have their original shapes and need to
+be broadcast such that the leading dimension is indexed by the items in the
+graph for which messages are computed (usually edges; except when convolving
+from nodes to context). In the end, values have to be pooled from there into a
+Tensor with a leading dimension indexed by receivers, see `pool_to_receiver`.
 
 
+
  

 
 
Args

 
 

 
-`sender_node_input`
+sender_node_input
 
 
-The input Tensor from the sender NodeSet, or `None`.
-If self.takes_sender_node_input is `False`, this arg will be `None`.
-(If it is `True`, that depends on how this layer gets called.)
+The input Tensor from the sender NodeSet, or None.
+If self.takes_sender_node_input is False, this arg will be None.
+(If it is True, that depends on how this layer gets called.)
 See also broadcast_from_sender_node.
 
 

 
-`sender_edge_input`
+sender_edge_input
 
 
-The input Tensor from the sender EdgeSet, or `None`.
-If self.takes_sender_edge_input is `False`, this arg will be `None`.
-(If it is `True`, it depends on how this layer gets called.)
+The input Tensor from the sender EdgeSet, or None.
+If self.takes_sender_edge_input is False, this arg will be None.
+(If it is True, it depends on how this layer gets called.)
 If present, this Tensor is already indexed by the items for which
 messages are computed.
 
 

 
-`receiver_input`
+receiver_input
 
 
 The input Tensor from the receiver NodeSet or Context,
-or None. If self.takes_receiver_input is `False`, this arg will be
-`None`. (If it is `True`, it depends on how this layer gets called.)
+or None. If self.takes_receiver_input is False, this arg will be
+None. (If it is True, it depends on how this layer gets called.)
 See broadcast_from_receiver.
 
 

 
-`broadcast_from_sender_node`
+broadcast_from_sender_node
 
 
 A function that broadcasts a Tensor indexed
@@ -311,25 +293,25 @@ messages are computed.
 
 

 
-`broadcast_from_receiver`
+broadcast_from_receiver
 
 
-Call this as `broadcast_from_receiver(value)`
+Call this as broadcast_from_receiver(value)
 to broadcast a Tensor indexed like receiver_input to a Tensor indexed
 by the items for which messages are computed.
 
 

 
-`pool_to_receiver`
+pool_to_receiver
 
 
-Call this as `pool_to_receiver(value, reduce_type=...)`
+Call this as pool_to_receiver(value, reduce_type=...)
 to pool an item-indexed Tensor to a receiver-indexed tensor, using
 a reduce_type understood by tfgnn.pool(), such as "sum".
 
 

 
-`extra_receiver_ops`
+extra_receiver_ops
 
 
 The extra_receiver_ops passed to init, see there,
@@ -339,10 +321,10 @@ this argument, so subclass implementors not using it can omit it.
 
 

 
-`training`
+training
 
 
-The `training` boolean that was passed to Layer.call(). If true,
+The training boolean that was passed to Layer.call(). If true,
 the result is computed for training rather than inference. For example,
 calls to tf.nn.dropout() are usually conditioned on this flag.
 By contrast, calling another Keras layer (like tf.keras.layers.Dropout)
@@ -351,9 +333,8 @@ does not require forwarding this arg, Keras does that automatically.
 

 

 
-
-
 
+
  
@@ -365,8 +346,3 @@ result of the convolution for each receiver.
 

 
 
Returns

 
 

-
-
-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Broadcast.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Broadcast.md
index 1884cc66..26cad901 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Broadcast.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Broadcast.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.Broadcast
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Broadcasts a GraphTensor feature.
 
@@ -60,13 +53,14 @@ left unset to select
 tfgnn.HIDDEN_STATE.
 
 
+
  

 
 
Init args

 
 

 
-`tag`
+tag
 
 
 Can be set to one of tfgnn.SOURCE, tfgnn.TARGET or tfgnn.CONTEXT
@@ -74,25 +68,25 @@ to select the sender from which feature values are broadcast.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 If set, the feature will be broadcast to this edge set
-(or this sequence of edge sets) from the sender given by `tag`.
-Mutually exclusive with `node_set_name`.
+(or this sequence of edge sets) from the sender given by tag.
+Mutually exclusive with node_set_name.
 
 

 
-`node_set_name`
+node_set_name
 
 
 If set, the feature will be broadcast to this node set
 (or sequence of node sets). The sender must be selected as
-`tag=tfgn.CONTEXT`. Mutually exclusive with `edge_set_name`.
+tag=tfgn.CONTEXT. Mutually exclusive with edge_set_name.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of the feature to read. If unset (also in call),
@@ -102,20 +96,21 @@ the tfgnn.HIDDEN_STATE
 

 
 
+
  

 
 
Call args

 
 

 
-`graph`
+graph
 
 
 The scalar tfgnn.GraphTensor to read from.
 
 

 
-`tag`
+tag
 
 
 Same meaning as for init. Must be passed to init, or to call,
@@ -125,7 +120,7 @@ edge_set_name, node_set_name: Same meaning as for init. One of them must
 
 

 
-`feature_name`
+feature_name
 
 
 Same meaning as for init. If passed to both, the value must
@@ -148,32 +143,29 @@ requested receivers.
 
 

 
-
-
-
-
 
+
  
+
+
+

 
 
Attributes

 
 

 
-`feature_name`
+feature_name
 
 
 Returns the feature_name argument to init, or None if unset.
 
 

 
-`location`
+location
 
 
 Returns dict of kwarg to init with the node or edge set name.
 
 

 
-`tag`
+tag
 
 
 Returns the tag argument to init, or None if unset.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ContextUpdate.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ContextUpdate.md
index 42aa6d9d..18a9b1dc 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ContextUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ContextUpdate.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ContextUpdate
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A context update with input from node sets and/or edge sets.
 
@@ -26,49 +19,53 @@ A context update with input from node sets and/or edge sets.
 )
 
 
-
-
 
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  
 
-
+

 
 
Init args
  `node_set_inputs`   A dict
-`{node_set_name: node_set_input, ...}` of Keras layers that return values shaped
-like context features with information aggregated from the given edge set. They
-are run on the input graph tensor as `node_set_input(graph,
-node_set_name=node_set_name)`.  
 
-`edge_set_inputs`   A dict `{edge_set_name:
-edge_set_input, ...}` of Keras layers that return values shaped like context
-features with information aggregated from the given edge set. They are run on
-the input graph tensor as `edge_set_input(graph, edge_set_name=edge_set_name)`.
- 
  `next_state`   A Keras
-layer to compute the new node state from a tuple of inputs that contains, in
-this order:
+
  node_set_inputs   A
+dict {node_set_name: node_set_input, ...} of Keras layers that
+return values shaped like context features with information aggregated from the
+given edge set. They are run on the input graph tensor as
+node_set_input(graph, node_set_name=node_set_name).  
 edge_set_inputs   A dict
+{edge_set_name: edge_set_input, ...} of Keras layers that return
+values shaped like context features with information aggregated from the given
+edge set. They are run on the input graph tensor as edge_set_input(graph,
+edge_set_name=edge_set_name).  
 
+next_state   A Keras layer to
+compute the new node state from a tuple of inputs that contains, in this order:
 
--   the `context_input_feature` (see there),
--   a dict `{node_set_name: input}` with the results of `node_set_inputs`, in
-    which each result is a tensor or dict of tensors,
--   a dict `{edge_set_name: input}` with the results of `edge_set_inputs`,
+-   the context_input_feature (see there),
+-   a dict {node_set_name: input} with the results of
+    node_set_inputs, in which each result is a tensor or dict of
+    tensors,
+-   a dict {edge_set_name: input} with the results of edge_set_inputs,
     in which each result is a tensor or dict of tensors, if there are any.
     
     

     
-    `context_input_feature`
+    context_input_feature
     
     
     The feature name(s) of inputs from the context to
-    `next_state`, defaults to tfgnn.HIDDEN_STATE.
+    next_state, defaults to tfgnn.HIDDEN_STATE.
     If set to a single feature name, a single tensor is passed.
-    If set to `None` or an empty sequence, an empty dict is passed.
+    If set to None or an empty sequence, an empty dict is passed.
     Otherwise, a dict of tensors keyed by feature names is passed.
     
     

     

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/EdgeSetUpdate.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/EdgeSetUpdate.md
index 3a9cc172..b35aa19b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/EdgeSetUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/EdgeSetUpdate.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.EdgeSetUpdate
 
-[TOC]
-
 
 
-

 
 
Call result

-
-
+
+ View source
+on GitHub 
 
 Computes the new state of an EdgeSet from select input features.
 
@@ -27,36 +20,38 @@ Computes the new state of an EdgeSet from select input features.
 )
 
 
-
-
 
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Init args

 
 

 
-`next_state`
+next_state
 
 
 The NextState layer to apply.
 
 

 
-`edge_input_feature`
+edge_input_feature
 
 
 The feature name(s) of inputs from the edge set to
-`next_state`, defaults to tfgnn.HIDDEN_STATE.
+next_state, defaults to tfgnn.HIDDEN_STATE.
 If set to a single feature name, a single tensor is passed.
-If set to `None` or an empty sequence, an empty dict is passed.
+If set to None or an empty sequence, an empty dict is passed.
 Otherwise, a dict of tensors keyed by feature names is passed.
 
 

 
-`node_input_tags`
+node_input_tags
 
 
 The incident nodes of each edge whose states are used
@@ -65,23 +60,23 @@ by default).
 
 

 
-`node_input_feature`
+node_input_feature
 
 
 The feature name of the input from node sets to
-`next_state`, defaults to tfgnn.HIDDEN_STATE.
-Setting this to `None` passes an empty dict of node inputs.
+next_state, defaults to tfgnn.HIDDEN_STATE.
+Setting this to None passes an empty dict of node inputs.
 This class supports only a single input feature from nodes. For more
 complex settings, you need to write your own, or start a design discussion
 about a node_input_map from tags to the respective features for each.
 
 

 
-`context_input_feature`
+context_input_feature
 
 
 The feature name(s) of inputs from the context to
-`next_state`. Defaults to `None`, which passes an empty dict.
+next_state. Defaults to None, which passes an empty dict.
 If set to a single feature name, a single tensor is passed.
 Otherwise, a dict of tensors keyed by feature names is passed.
 To pass the default state tensor of the context, set this to
@@ -91,6 +86,7 @@ To pass the default state tensor of the context, set this to
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/GraphUpdate.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/GraphUpdate.md
index 1249d28d..c10b986c 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/GraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/GraphUpdate.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.GraphUpdate
 
-[TOC]
-
 
 
-

 
 
Call returns

-
-
+
+ View source
+on GitHub 
 
 Applies one round of updates to EdgeSets, NodeSets and Context.
 
@@ -26,28 +19,30 @@ Applies one round of updates to EdgeSets, NodeSets and Context.
 )
 
 
-
-
 
 
-The updates of EdgeSets, NodeSets and Context can either be passed as
-init arguments, or constructed later by passing a deferred_init_callback,
-which allows advanced users to adjust the updates to the GraphTensorSpec
-of the input (which EdgeSets and NodeSets even exist).
+The updates of EdgeSets, NodeSets and Context can either be passed as init
+arguments, or constructed later by passing a deferred_init_callback, which
+allows advanced users to adjust the updates to the GraphTensorSpec of the input
+(which EdgeSets and NodeSets even exist).
+
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
 
 
+
  

 
 
Init args

 
 

 
-`edge_sets`
+edge_sets
 
 
-A dict `{edge_set_name: edge_set_update, ...}` of EdgeSetUpdate
+A dict {edge_set_name: edge_set_update, ...} of EdgeSetUpdate
 layers (or custom reimplementations). They are run on the input graph
-tensor as `edge_set_update(graph, edge_set_name=edge_set_name)`.
+tensor as edge_set_update(graph, edge_set_name=edge_set_name).
 Their results are merged into the feature map of the respective edge set.
 This argument can be omitted, which is common in models with node set
 updates that use convolutions (i.e., read from adjacent nodes without
@@ -55,19 +50,19 @@ computing explicit edge states).
 
 

 
-`node_sets`
+node_sets
 
 
-A dict `{node_set_name: node_set_update, ...}` of NodeSetUpdate
+A dict {node_set_name: node_set_update, ...} of NodeSetUpdate
 layers (or custom reimplementations). They are run on the graph tensor
 with edge set updates (if any) as
-`node_set_update(graph, node_set_name=node_set_name)`,
+node_set_update(graph, node_set_name=node_set_name),
 Their results are merged into the feature map of the respective node set.
 This argument can be omitted (but that is uncommon).
 
 

 
-`context`
+context
 
 
 A ContextUpdate that is run on the graph tensor with edge set and
@@ -77,7 +72,7 @@ without a context state.
 
 

 
-`deferred_init_callback`
+deferred_init_callback
 
 
 Can be set to a function that accepts a
@@ -93,6 +88,7 @@ be saved.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ItemDropout.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ItemDropout.md
index afd222c9..0c3dbd0c 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ItemDropout.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ItemDropout.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ItemDropout
 
-[TOC]
-
 
 
-

 
 
Call result

-
-
+
+ View source
+on GitHub 
 
 Dropout of feature values for entire edges, nodes or components.
 
@@ -27,39 +20,44 @@ This Layer class wraps `tf.keras.layers.Dropout` to perform edge dropout or node
 dropout (or "component dropout", which is rarely useful) on Tensors shaped like
 features of a **scalar** GraphTensor.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Init args

 
 

 
-`rate`
+rate
 
 
-The dropout rate, forwarded to `tf.keras.layers.Dropout`.
+The dropout rate, forwarded to tf.keras.layers.Dropout.
 
 

 
-`seed`
+seed
 
 
-The random seed, forwarded `tf.keras.layers.Dropout`.
+The random seed, forwarded tf.keras.layers.Dropout.
 
 

 

 
 
+
  
@@ -67,15 +65,16 @@ GraphTensor. Across calls, all inputs must have the same known rank.
 

 
 
Call args

 
 

 
-`x`
+x
 
 
-A float Tensor of shape `[num_items, *feature_dims]`. This is the shape
+A float Tensor of shape [num_items, *feature_dims]. This is the shape
 of node features or edge features (or context features) in a *scalar**
 GraphTensor. Across calls, all inputs must have the same known rank.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MakeEmptyFeature.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MakeEmptyFeature.md
index b5257840..a2c13f00 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MakeEmptyFeature.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MakeEmptyFeature.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.MakeEmptyFeature
 
-[TOC]
-
 
 
-

 
 
Call returns

 

 
-A Tensor `y` with the same shape and dtype as the input `x`.
-In non-training mode, the output is the same as the input: `y == x`.
-In training mode, each row `y[i]` is either zeros (with probability `rate`)
-or a scaled-up copy of the input row: `y[i] = x[i] * 1./(1-rate)`.
+A Tensor y with the same shape and dtype as the input x.
+In non-training mode, the output is the same as the input: y == x.
+In training mode, each row y[i] is either zeros (with probability rate)
+or a scaled-up copy of the input row: y[i] = x[i] * 1./(1-rate).
 This is similar to ordinary dropout, except all or none of the feature
 values for each item are dropped out.
 

-
-
+
+ View source
+on GitHub 
 
 Returns an empty feature with a shape that fits the input graph piece.
 
@@ -21,25 +14,27 @@ Returns an empty feature with a shape that fits the input graph piece.
 )
 
 
-
-
 
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Init args

 
 

 
-`dtype`
+dtype
 
 
 the tf.DType to use for the result, defaults to tf.float32.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Other arguments for the tf.keras.layers.Layer base class.
@@ -48,13 +43,14 @@ Other arguments for the tf.keras.layers.Layer base class.
 

 
 
+
  

 
 
Call args

 
 

 
-`graph_piece`
+graph_piece
 
 
 a Context, NodeSet or EdgeSet from a GraphTensor.
@@ -63,6 +59,7 @@ a Context, NodeSet or EdgeSet from a GraphTensor.
 

 
 
+
  
@@ -80,6 +77,7 @@ GraphTensor, the result is a Tensor of shape [graph_piece.total_size, 0].
 

 
 
Call returns

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MapFeatures.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MapFeatures.md
index 2b1f2b90..13700a35 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MapFeatures.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/MapFeatures.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.MapFeatures
 
-[TOC]
-
 
 
-

 
 
TPU compatibility

-
-
+
+ View source
+on GitHub 
 
 Transforms features on a GraphTensor by user-defined callbacks.
 
@@ -30,10 +23,10 @@ Transforms features on a GraphTensor by user-defined callbacks.
 
 
 This layer transforms the feature maps of graph pieces (that is, EdgeSets,
-NodeSets, or the Context) by applying Keras Models to them. Those Models
-are built by user-supplied callbacks that receive a KerasTensor for the
-graph piece as input and return a dict of output features computed with
-the Keras functional API, see https://tensorflow.org/guide/keras/functional.
+NodeSets, or the Context) by applying Keras Models to them. Those Models are
+built by user-supplied callbacks that receive a KerasTensor for the graph piece
+as input and return a dict of output features computed with the Keras functional
+API, see https://tensorflow.org/guide/keras/functional.
 
 Auxiliary graph pieces (e.g., for
 tfgnn.keras.layers.StructuredReadout)
@@ -42,8 +35,6 @@ are skipped, unless explicitly requested via `allowed_aux_node_sets_pattern` or
 
 #### Examples:
 
-
-
 ```python
 # Hashes edge features called "id", leaves others unchanged:
 def edge_sets_fn(edge_set, *, edge_set_name):
@@ -81,11 +72,11 @@ graph = tfgnn.keras.layers.MapFeatures(
 )(graph)
 ```
 
-When this layer is called on a GraphTensor, it transforms the feature map
-of each graph piece with the model built by the respective callbacks.
-The very first call to this layer triggers building the models. Subsequent
-calls to this layer do not use the callbacks again, but check that their
-input does not have more graph pieces or features than seen by the callbacks:
+When this layer is called on a GraphTensor, it transforms the feature map of
+each graph piece with the model built by the respective callbacks. The very
+first call to this layer triggers building the models. Subsequent calls to this
+layer do not use the callbacks again, but check that their input does not have
+more graph pieces or features than seen by the callbacks:
 
 *   It is an error to call with a node set or edge set that was not present in
     the first call. (After the first call, it is too late to initialize another
@@ -102,34 +93,33 @@ input does not have more graph pieces or features than seen by the callbacks:
 More details on the callbacks:
 
 The model-building callbacks are passed as arguments when initializing this
-layer (see "Init args" below). Each callback is invoked as
-`fn(graph_piece, **kwargs)` where
+layer (see "Init args" below). Each callback is invoked as `fn(graph_piece,
+**kwargs)` where
 
-  * `graph_piece` is a KerasTensor for the EdgeSet, NodeSet or Context
-    that is being transformed. It provides access to the input features.
-  * the keyword argument (if any) is
-      * `edge_set_name=...` when transforming the features of that EdgeSet,
-      * `node_set_name=...` when transforming the features of that NodeSet,
-      * absent when transforming the features of the Context.
+*   `graph_piece` is a KerasTensor for the EdgeSet, NodeSet or Context that is
+    being transformed. It provides access to the input features.
+*   the keyword argument (if any) is
+    *   `edge_set_name=...` when transforming the features of that EdgeSet,
+    *   `node_set_name=...` when transforming the features of that NodeSet,
+    *   absent when transforming the features of the Context.
 
 The output of the callbacks can take one of the following forms:
 
-  * A returned dict of feature values is used as the new feature map of
-    the respective graph piece in this layer's output. Returning the
-    empty dict `{}` is allowed and results in an empty feature map.
-  * A returned feature `value` not wrapped in a dict is a shorthand for
-    `{tfgnn.HIDDEN_STATE: value}`, to simplify the set-up of initial
-    states.
-  * Returning `None` as the callback's result indicates to leave this graph
+*   A returned dict of feature values is used as the new feature map of the
+    respective graph piece in this layer's output. Returning the empty dict `{}`
+    is allowed and results in an empty feature map.
+*   A returned feature `value` not wrapped in a dict is a shorthand for
+    `{tfgnn.HIDDEN_STATE: value}`, to simplify the set-up of initial states.
+*   Returning `None` as the callback's result indicates to leave this graph
     piece alone and not even validate that subsequent inputs have the same
     features.
 
 The output values are required to
 
-  * have the correct shape for a feature on the respective piece of the
+*   have the correct shape for a feature on the respective piece of the
     GraphTensor;
-  * depend on the input, so that the Keras functional API can use them
-    as Model outputs.
+*   depend on the input, so that the Keras functional API can use them as Model
+    outputs.
 
 This happens naturally for outputs of transformed input features. Outputs
 created from scratch still need to depend on the input for its size. The helper
@@ -144,75 +134,80 @@ KerasTensors.
 Weight sharing between the transformation of different graph pieces is possible
 by sharing the Keras objects between the respective callback invocations.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 WARNING: Weight sharing fails in `tf.keras.models.load_model()` with an error
 message on weights missing from the checkpoint. (Most users don't need to
 re-load their models this way.)
 
 
+
  

 
 
Init args

 
 

 
-`context_fn`
+context_fn
 
 
 A callback to build a Keras model for transforming context
-features. It will be called as `output = context_fn(g.context)`.
-Leaving this at the default `None` is equivalent to returning `None`.
+features. It will be called as output = context_fn(g.context).
+Leaving this at the default None is equivalent to returning None.
 
 

 
-`node_sets_fn`
+node_sets_fn
 
 
 A callback to build a Keras model for transforming node set
 features. It will be called for every node sets as
-`node_sets_fn(g.node_sets[node_set_name], node_set_name=node_set_name)`.
-Leaving this at the default `None` is equivalent to returning `None`
+node_sets_fn(g.node_sets[node_set_name], node_set_name=node_set_name).
+Leaving this at the default None is equivalent to returning None
 for every node set.
 
 

 
-`edge_sets_fn`
+edge_sets_fn
 
 
 A callback to build a Keras model for transforming edge set
 features. It will be called for every edge sets as
-`edge_sets_fn(g.edge_sets[edge_set_name], edge_set_name=edge_set_name)`.
-Leaving this at the default `None` is equivalent to returning `None`
+edge_sets_fn(g.edge_sets[edge_set_name], edge_set_name=edge_set_name).
+Leaving this at the default None is equivalent to returning None
 for every edge set.
 
 

 
-`allowed_aux_node_sets_pattern`
+allowed_aux_node_sets_pattern
 
 
-If set, `node_sets_fn` is also invoked for
+If set, node_sets_fn is also invoked for
 those auxiliary node sets that match this pattern, according to Python's
-`re.fullmatch(pattern, node_set_name)`.
+re.fullmatch(pattern, node_set_name).
 
 

 
-`allowed_aux_edge_sets_pattern`
+allowed_aux_edge_sets_pattern
 
 
-If set, `edge_sets_fn` is also invoked for
+If set, edge_sets_fn is also invoked for
 those auxiliary edge sets that match this pattern, according to Python's
-`re.fullmatch(pattern, edge_set_name)`.
+re.fullmatch(pattern, edge_set_name).
 
 

 

 
 
+
  

 
 
Call args

 
 

 
-`graph`
+graph
 
 
 A GraphTensor. The very first call triggers the building of
@@ -223,6 +218,7 @@ taken from the GraphTensorSpec of the first input.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NextStateFromConcat.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NextStateFromConcat.md
index 2b5d3d2f..b409eef2 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NextStateFromConcat.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NextStateFromConcat.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.NextStateFromConcat
 
-[TOC]
-
 
 
-

 
 
Call returns

-
-
+
+ View source
+on GitHub 
 
 Computes a new state by concatenating inputs and applying a Keras Layer.
 
@@ -21,21 +14,23 @@ Computes a new state by concatenating inputs and applying a Keras Layer.
 )
 
 
-
-
 
 
 This layer flattens all inputs into a list (forgetting their origin),
 concatenates them and sends them through a user-supplied feed-forward network.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  

 
 
Init args

 
 

 
-`transformation`
+transformation
 
 
 Required. A Keras Layer to transform the combined inputs
@@ -45,6 +40,7 @@ into the new state.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NodeSetUpdate.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NodeSetUpdate.md
index 14972a35..ac73bdbc 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NodeSetUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/NodeSetUpdate.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.NodeSetUpdate
 
-[TOC]
-
 
 
-

 
 
Call returns

-
-
+
+ View source
+on GitHub 
 
 A node state update with input from convolutions or other edge set inputs.
 
@@ -26,46 +19,50 @@ A node state update with input from convolutions or other edge set inputs.
 )
 
 
-
-
 
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  
 
-
+

 
 
Init args
  `edge_set_inputs`   A dict
-`{edge_set_name: edge_set_input, ...}` of Keras layers (such as convolutions)
-that return values shaped like node features with information aggregated from
-the given edge set. They are run in parallel on the input graph tensor as
-`edge_set_input(graph, edge_set_name=edge_set_name)`.  
 
-`next_state`   A Keras layer to compute the new
-node state from a tuple of inputs that contains, in this order:
+
  edge_set_inputs   A
+dict {edge_set_name: edge_set_input, ...} of Keras layers (such as
+convolutions) that return values shaped like node features with information
+aggregated from the given edge set. They are run in parallel on the input graph
+tensor as edge_set_input(graph, edge_set_name=edge_set_name). 
  next_state   A Keras
+layer to compute the new node state from a tuple of inputs that contains, in
+this order:
 
--   the `node_input_feature` (see there),
--   a dict `{edge_set_name: input}` with the results of `edge_set_inputs`, in
-    which each result is a tensor or dict of tensors,
--   if context_input_feature is not `None`, those feature(s).
+-   the node_input_feature (see there),
+-   a dict {edge_set_name: input} with the results of
+    edge_set_inputs, in which each result is a tensor or dict of
+    tensors,
+-   if context_input_feature is not None, those feature(s).
     
     

     
-    `node_input_feature`
+    node_input_feature
     
     
     The feature name(s) of inputs from the node set to
-    `next_state`, defaults to tfgnn.HIDDEN_STATE.
+    next_state, defaults to tfgnn.HIDDEN_STATE.
     If set to a single feature name, a single tensor is passed.
-    If set to `None` or an empty sequence, an empty dict is passed.
+    If set to None or an empty sequence, an empty dict is passed.
     Otherwise, a dict of tensors keyed by feature names is passed.
     
     

     
-    `context_input_feature`
+    context_input_feature
     
     
     The feature name(s) of inputs from the context to
-    `next_state`. Defaults to `None`, which passes an empty dict.
+    next_state. Defaults to None, which passes an empty dict.
     If set to a single feature name, a single tensor is passed.
     Otherwise, a dict of tensors keyed by feature names is passed.
     To pass the default state tensor of the context, set this to
@@ -75,6 +72,7 @@ node state from a tuple of inputs that contains, in this order:
     

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/PadToTotalSizes.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/PadToTotalSizes.md
index b678656a..9c9acb53 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/PadToTotalSizes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/PadToTotalSizes.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.PadToTotalSizes
 
-[TOC]
-
 
 
-

 
 
Call result

-
-
+
+ View source
+on GitHub 
 
 Applies tfgnn.pad_to_total_sizes() to a GraphTensor.
 
@@ -24,15 +17,14 @@ Applies tfgnn.pad_to_total_sizes() to a GraphTensor.
 )
 
 
-
-
 
 
 This Keras layer maps a GraphTensor to a GraphTensor by calling
-tfgnn.pad_to_total_sizes() with the additional arguments, notably
-`sizes_constraints`, passed at initialization time. See that function
-for detailed documentation.
-
-Serialization to a Keras model config requires the `sizes_constraints` to
-contain Python integers or eager Tensors, not symbolic Tensors.
-
+tfgnn.pad_to_total_sizes()
+with the additional arguments, notably `sizes_constraints`, passed at
+initialization time. See that function for detailed documentation.
+
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`. Serialization to a
+Keras model config requires the `sizes_constraints` to contain Python integers
+or eager Tensors, not symbolic Tensors.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseExample.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseExample.md
index 08709382..d2856137 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseExample.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseExample.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ParseExample
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Applies tfgnn.parse_example(graph_tensor_spec, _) to a batch of strings.
 
@@ -22,8 +15,7 @@ Applies tfgnn.parse_example(graph_tensor_spec, _) to a batch of strings.
 )
 
 
-
-
 
 
-
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseSingleExample.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseSingleExample.md
index afb0e7ad..37bcaed5 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseSingleExample.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ParseSingleExample.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ParseSingleExample
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Applies tfgnn.parse_single_example(graph_tensor_spec, _).
 
@@ -22,8 +15,7 @@ Applies tfgnn.parse_single_example(graph_tensor_spec, _).
 )
 
 
-
-
 
 
-
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Pool.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Pool.md
index 9f1de463..1bbf5565 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Pool.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Pool.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.Pool
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Pools a GraphTensor feature.
 
@@ -29,8 +22,8 @@ Pools a GraphTensor feature.
 
 
 
-This layer accepts a complete GraphTensor and returns a tensor with the
-result of pooling some feature.
+This layer accepts a complete GraphTensor and returns a tensor with the result
+of pooling some feature.
 
 There are two kinds of pooling that this layer can be used for:
 
@@ -64,13 +57,14 @@ The feature name can be left unset to select
 tfgnn.HIDDEN_STATE.
 
 
+
  

 
 
Init args

 
 

 
-`tag`
+tag
 
 
 Can be set to one of tfgnn.SOURCE, tfgnn.TARGET or tfgnn.CONTEXT
@@ -78,32 +72,32 @@ to select the receiver.
 
 

 
-`reduce_type`
+reduce_type
 
 
-Can be set to any `reduce_type` understood by tfgnn.pool().
+Can be set to any reduce_type understood by tfgnn.pool().
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 If set, the feature will be pooled from this edge set
-(or this sequence of edge sets) to the receiver given by `tag`.
-Mutually exclusive with `node_set_name`.
+(or this sequence of edge sets) to the receiver given by tag.
+Mutually exclusive with node_set_name.
 
 

 
-`node_set_name`
+node_set_name
 
 
 If set, the feature will be pooled from this node set
 (or sequence of node sets). The receiver must be selected as
-`tag=tfgnn.CONTEXT`. Mutually exclusive with `edge_set_name`.
+tag=tfgnn.CONTEXT. Mutually exclusive with edge_set_name.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of the feature to read. If unset (also in call),
@@ -113,20 +107,21 @@ the tfgnn.HIDDEN_STATE
 

 
 
+
  

 
 
Call args

 
 

 
-`graph`
+graph
 
 
 The scalar tfgnn.GraphTensor to read from.
 
 

 
-`reduce_type`
+reduce_type
 
 
 Same meaning as for init. Must be passed to init, or to call,
@@ -134,7 +129,7 @@ or to both (with the same value).
 
 

 
-`tag`
+tag
 
 
 Same meaning as for init. Must be passed to init, or to call,
@@ -144,7 +139,7 @@ edge_set_name, node_set_name: Same meaning as for init. One of them must
 
 

 
-`feature_name`
+feature_name
 
 
 Same meaning as for init. If passed to both, the value must
@@ -154,6 +149,7 @@ be the same. If passed to neither, 
 
 
+
  
@@ -165,39 +161,36 @@ A tensor with the pooled feature value.
 
 

 
 
Returns

 
-
-
-
-
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ResidualNextState.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ResidualNextState.md
index 2c137c4a..4626c9c8 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ResidualNextState.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ResidualNextState.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ResidualNextState
 
-[TOC]
-
 
 
-

 
 
Attributes

 
 

 
-`feature_name`
+feature_name
 
 
 Returns the feature_name argument to init, or None if unset.
 
 

 
-`location`
+location
 
 
 Returns dict of kwarg to init with the node or edge set name.
 
 

 
-`reduce_type`
+reduce_type
 
 
 Returns the reduce_type argument to init, or None if unset.
 
 

 
-`tag`
+tag
 
 
 Returns the tag argument to init, or None if unset.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Readout.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Readout.md
index d08783bd..df4d238d 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Readout.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/Readout.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.Readout
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Reads a feature out of a GraphTensor.
 
@@ -36,14 +29,14 @@ require a Keras Layer and do not allow subscrpting syntax like
 `graph_tensor.node_sets["user"]["name"]`.
 
 A location in the graph is selected by setting exactly one of the keyword
-arguments `edge_set_name=...`, `node_set_name=...` or `from_context=True`.
-From there, the keyword argument `feature_name=...` selects the feature.
+arguments `edge_set_name=...`, `node_set_name=...` or `from_context=True`. From
+there, the keyword argument `feature_name=...` selects the feature.
 
-Both the initialization of and the call to this layer accept arguments to
-select the feature location and the feature name. The call arguments take
-effect for that call only and can supply missing values, but they are not
-allowed to contradict initialization arguments. The feature name can be left
-unset to select tfgnn.HIDDEN_STATE.
+Both the initialization of and the call to this layer accept arguments to select
+the feature location and the feature name. The call arguments take effect for
+that call only and can supply missing values, but they are not allowed to
+contradict initialization arguments. The feature name can be left unset to
+select tfgnn.HIDDEN_STATE.
 
 #### For example:
 
@@ -64,13 +57,14 @@ previously handled by
 tfgnn.keras.layers.ReadoutFirstNode.
 
 
+
  

 
 
Init args

 
 

 
-`edge_set_name`
+edge_set_name
 
 
 If set, the feature will be read from this edge set.
@@ -78,7 +72,7 @@ Mutually exclusive with node_set_name and from_context.
 
 

 
-`node_set_name`
+node_set_name
 
 
 If set, the feature will be read from this node set.
@@ -86,7 +80,7 @@ Mutually exclusive with edge_set_name and from_context.
 
 

 
-`from_context`
+from_context
 
 
 If true, the feature will be read from the context.
@@ -94,7 +88,7 @@ Mutually exclusive with edge_set_name and node_set_name.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of the feature to read. If unset (also in call),
@@ -104,13 +98,14 @@ tfgnn.HIDDEN_STATE will be read.
 

 
 
+
  

 
 
Call args

 
 

 
-`graph`
+graph
 
 
 The GraphTensor to read from.
@@ -119,7 +114,7 @@ edge_set_name, node_set_name, from_context: Same meaning as for init. One of
 
 

 
-`feature_name`
+feature_name
 
 
 Same meaning as for init. If passed to both, the value must
@@ -129,6 +124,7 @@ be the same. If passed to neither, tfgnn.HIDDEN_STATE is used.
 

 
 
+
  
@@ -140,28 +136,22 @@ The tensor with the selected feature.
 
 

 
 
Returns

 
-
-
-
-
 
+
  
 
-

 
 
Attributes
  `feature_name`   Returns the
-feature_name argument to init, or None if unset.  
 
-`location`   Returns a dict with the kwarg to init
-that selected the feature location.
+
  feature_name   Returns
+the feature_name argument to init, or None if unset.  
 
+location   Returns a dict with the
+kwarg to init that selected the feature location.
 
-The result contains the keyword argument and value passed to `__init__()`
+The result contains the keyword argument and value passed to **init**()
 that selects the location from which the layer's output feature is read,
- that is, one of `edge_set_name=...`, `node_set_name=...` or
-`from_context=True`. If none of these has been set, the result is
+ that is, one of edge_set_name=..., node_set_name=... or
+from_context=True. If none of these has been set, the result is
 empty, and one of them must be set at call time.
 
 

 

-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutFirstNode.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutFirstNode.md
index 7f551c39..2e1e6382 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutFirstNode.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutFirstNode.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ReadoutFirstNode
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Reads a feature from the first node of each graph component.
 
@@ -26,15 +19,15 @@ Reads a feature from the first node of each graph component.
 
 
 
-Given a particular node set (identified by `node_set_name`), this layer
-will gather the given feature from the first node of each graph component.
+Given a particular node set (identified by `node_set_name`), this layer will
+gather the given feature from the first node of each graph component.
 
 This is often used for rooted graphs created by sampling around the
 neighborhoods of seed nodes in a large graph: by convention, each seed node is
-the first node of its component in the respective node set, and this layer
-reads out the information it has accumulated there. (In other node sets, the
-first node may be arbitrary -- or nonexistant, in which case this operation
-must not be used and may raise an error at runtime.)
+the first node of its component in the respective node set, and this layer reads
+out the information it has accumulated there. (In other node sets, the first
+node may be arbitrary -- or nonexistent, in which case this operation must not
+be used and may raise an error at runtime.)
 
 This implicit convention is inflexible and hard to validate. New models are
 encouraged to use
@@ -44,20 +37,21 @@ sampled dataset itself; if not, it can be added after the fact by
 tfgnn.keras.layers.AddReadoutFromFirstNode.
 
 
+
  

 
 
Init args

 
 

 
-`node_set_name`
+node_set_name
 
 
 If set, the feature will be read from this node set.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of the feature to read. If unset (also in call),
@@ -67,20 +61,21 @@ tfgnn.HIDDEN_STATE will be read.
 

 
 
+
  

 
 
Call args

 
 

 
-`graph`
+graph
 
 
 The scalar GraphTensor to read from.
 
 

 
-`node_set_name`
+node_set_name
 
 
 Same meaning as for init. Must be passed to init, or to call,
@@ -88,7 +83,7 @@ or to both (with the same value).
 
 

 
-`feature_name`
+feature_name
 
 
 Same meaning as for init. If passed to both, the value must
@@ -98,6 +93,7 @@ be the same. If passed to neither, tfgnn.HIDDEN_STATE is used.
 

 
 
+
  
@@ -110,23 +106,17 @@ context feature.
 
 

 
 
Returns

 
-
-
-
-
 
+
  
 
-

 
 
Attributes
  `feature_name`   Returns the
-feature_name argument to init, or None if unset.  
 
-`location`   Returns a dict with the kwarg to init
-that selected the feature location.
+
  feature_name   Returns
+the feature_name argument to init, or None if unset.  
 
+location   Returns a dict with the
+kwarg to init that selected the feature location.
 
 
 

 

-
-
-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutNamedIntoFeature.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutNamedIntoFeature.md
index 48797d43..b57c41a7 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutNamedIntoFeature.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ReadoutNamedIntoFeature.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.ReadoutNamedIntoFeature
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Reads out a feature value from select nodes (or edges) in a graph.
 
@@ -61,16 +54,16 @@ training labels out of the input graph seen my the model (see
 
 

 
-`key`
+key
 
 
 A string key to select between possibly multiple named readouts
-(such as `"source"` and `"target"` for link prediction). Can be fixed
+(such as "source" and "target" for link prediction). Can be fixed
 in init, or selected for each call.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of a feature to read out from, as with
@@ -78,42 +71,42 @@ The name of a feature to read out from, as with
 
 

 
-`new_feature_name`
+new_feature_name
 
 
-The name of the feature to add to `readout_node_set`
-for storing the readout result. If unset, defaults to `feature_name`.
-It is an error if the added feature already exists on `readout_node_set`
-in the input `graph`, unless `overwrite=True` is set.
+The name of the feature to add to readout_node_set
+for storing the readout result. If unset, defaults to feature_name.
+It is an error if the added feature already exists on readout_node_set
+in the input graph, unless overwrite=True is set.
 
 

 
-`remove_input_feature`
+remove_input_feature
 
 
-If set, the given `feature_name` is removed from the
+If set, the given feature_name is removed from the
 node (or edge) set(s) that contain the value to be read out in the input
-`GraphTensor`.
+GraphTensor.
 
 

 
-`overwrite`
+overwrite
 
 
 If set, allows overwriting a potentially already existing
-feature `graph.node_sets[readout_node_set][new_feature_name]`.
+feature graph.node_sets[readout_node_set][new_feature_name].
 
 

 
-`readout_node_set`
+readout_node_set
 
 
-A string, defaults to `"_readout"`. This is used as the
+A string, defaults to "_readout". This is used as the
 name for the readout node set and as a name prefix for its edge sets.
 
 

 
-`validate`
+validate
 
 
 Setting this to false disables the validity checks of
@@ -132,14 +125,14 @@ on structurally unchanged GraphTensors.
 
 

 
-`graph`
+graph
 
 
-The scalar `GraphTensor` to read from.
+The scalar GraphTensor to read from.
 
 

 
-`key`
+key
 
 
 Same meaning as for init. Must be passed to init, or to call,
@@ -155,9 +148,9 @@ or to both (with the same value).
 
Returns

 

 
-A `GraphTensor` like `graph`, with the readout result stored as
-`.node_sets[readout_node_set][new_feature_name]` and possibly the
-readout input feature removed (see `remove_input_feature`).
+A GraphTensor like graph, with the readout result stored as
+.node_sets[readout_node_set][new_feature_name] and possibly the
+readout input feature removed (see remove_input_feature).
 
 

-
-
+
+ View source
+on GitHub 
 
 Updates a state with a residual block.
 
@@ -38,14 +31,18 @@ If the initial state of the graph piece that is being updated has size 0, the
 skip connection is omitted. This avoids the need to special-case, say, latent
 node sets in modeling code applied to different node sets.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/learn_fit_or_skip_size_constraints.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/learn_fit_or_skip_size_constraints.md
index 31963e69..f272b479 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/learn_fit_or_skip_size_constraints.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/learn_fit_or_skip_size_constraints.md
@@ -1,17 +1,10 @@
 # tfgnn.learn_fit_or_skip_size_constraints
 
-[TOC]
-
 
 
-

 
 
Init args

 
 

 
-`residual_block`
+residual_block
 
 
 Required. A Keras Layer to transform the concatenation
@@ -53,18 +50,18 @@ of all inputs into a delta that gets added to the state.
 
 

 
-`activation`
+activation
 
 
 An activation function (none by default), as understood by
-`tf.keras.layers.Activation`. This activation function is applied after
+tf.keras.layers.Activation. This activation function is applied after
 the residual block and the addition. If using this, typically the
 residual block does not have an activation function on its last layer,
 or vice versa.
 
 

 
-`skip_connection_feature_name`
+skip_connection_feature_name
 
 
 Controls which input from the updated graph
@@ -77,6 +74,7 @@ defaults to tfgnn.HIDDEN_STATE
 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SimpleConv.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SimpleConv.md
index aa1d784f..ff4cc4a6 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SimpleConv.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SimpleConv.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.SimpleConv
 
-[TOC]
-
 
 
-

 
 
Call returns

-
-
+
+ View source
+on GitHub 
 
 A convolution layer that applies a passed-in message_fn.
 
@@ -59,14 +52,19 @@ graph = tfgnn.keras.layers.GraphUpdate(
 )(graph)
 ```
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`, provided that the
+layers used in `message_fn` also support this. (Keras core layers do.)
+
 
+
  

 
 
Init args

 
 

 
-`message_fn`
+message_fn
 
 
 A Keras layer that computes the individual messages from the
@@ -74,17 +72,17 @@ combined input features (see combine_type).
 
 

 
-`reduce_type`
+reduce_type
 
 
 Specifies how to pool the messages to receivers. Defaults to
-`"sum"`, can be any reduce_type understood by `tfgnn.pool()`, including
-concatenations like `"sum|max"` (but mind the increased dimension of the
+"sum", can be any reduce_type understood by tfgnn.pool(), including
+concatenations like "sum|max" (but mind the increased dimension of the
 result and the growing number of model weights in the next-state layer).
 
 

 
-`combine_type`
+combine_type
 
 
 a string understood by tfgnn.combine_values(), to specify how
@@ -93,7 +91,7 @@ to "concat", which concatenates inputs along the last axis.
 
 

 
-`receiver_tag`
+receiver_tag
 
 
  one of tfgnn.SOURCE, tfgnn.TARGET or tfgnn.CONTEXT.
@@ -106,39 +104,40 @@ If left unset for init, the tag must be passed at call time.
 
 

 
-`receiver_feature`
+receiver_feature
 
 
 Can be set to override tfgnn.HIDDEN_STATE for use as
-the input feature from the receiver. Passing `None` disables input from
+the input feature from the receiver. Passing None disables input from
 the receiver.
 
 

 
-`sender_node_feature`
+sender_node_feature
 
 
 Can be set to override tfgnn.HIDDEN_STATE for use as
-the input feature from sender nodes. Passing `None` disables input from
+the input feature from sender nodes. Passing None disables input from
 the sender node.
-IMPORANT: Must be set to `None` for use with `receiver_tag=tfgnn.CONTEXT`
+IMPORTANT: Must be set to None for use with receiver_tag=tfgnn.CONTEXT
 on an edge set, or for pooling from edges without sender node states.
 
 

 
-`sender_edge_feature`
+sender_edge_feature
 
 
 Can be set to a feature name of the edge set to select
-it as an input feature. By default, this set to `None`, which disables
+it as an input feature. By default, this set to None, which disables
 this input.
-IMPORTANT: Must be set for use with `receiver_tag=tfgnn.CONTEXT` on an
+IMPORTANT: Must be set for use with receiver_tag=tfgnn.CONTEXT on an
 edge set.
 
 

 

 
 
+
  
@@ -152,11 +151,12 @@ pooled messages for each receiver.
 

 
 
Call returns

 
 
+
  
 
-
+

 
 
Args
  `receiver_tag`   one of
+
  receiver_tag   one of
 tfgnn.SOURCE,
 tfgnn.TARGET or
 tfgnn.CONTEXT. The results
@@ -168,45 +168,46 @@ of the edges. If set to
 tfgnn.CONTEXT, the layer
 can be called for an edge set or a node set and will aggregate results for
 context (per graph component). If left unset for init, the tag must be passed at
-call time.  
  `receiver_feature`
-  The name of the feature that is read from the receiver graph piece
-and passed as convolve(receiver_input=...).  
 
-`sender_node_feature`   The name of the
-feature that is read from the sender nodes, if any, and passed as
-convolve(sender_node_input=...). NOTICE this must be `None` for use with
-`receiver_tag=tfgnn.CONTEXT` on an edge set, or for pooling from edges without
-sender node states.  
 
-`sender_edge_feature`   The name of the
-feature that is read from the sender edges, if any, and passed as
-convolve(sender_edge_input=...). NOTICE this must not be `None` for use with
-`receiver_tag=tfgnn.CONTEXT` on an edge set.  
 
-`extra_receiver_ops`   A str-keyed
-dictionary of Python callables that are wrapped to bind some arguments and then
-passed on to `convolve()`. Sample usage: `extra_receiver_ops={"softmax":
-tfgnn.softmax}`. The values passed in this dict must be callable as follows,
-with two positional arguments:
+call time.  
 
+receiver_feature   The name
+of the feature that is read from the receiver graph piece and passed as
+convolve(receiver_input=...).  
 
+sender_node_feature   The
+name of the feature that is read from the sender nodes, if any, and passed as
+convolve(sender_node_input=...). NOTICE this must be None for use
+with receiver_tag=tfgnn.CONTEXT on an edge set, or for pooling from
+edges without sender node states.  
 
+sender_edge_feature   The
+name of the feature that is read from the sender edges, if any, and passed as
+convolve(sender_edge_input=...). NOTICE this must not be None for
+use with receiver_tag=tfgnn.CONTEXT on an edge set.  
 extra_receiver_ops   A
+str-keyed dictionary of Python callables that are wrapped to bind some arguments
+and then passed on to convolve(). Sample usage:
+extra_receiver_ops={"softmax": tfgnn.softmax}. The values passed in
+this dict must be callable as follows, with two positional arguments:
 
 ```python
 f(graph, receiver_tag, node_set_name=..., feature_value=..., ...)
 f(graph, receiver_tag, edge_set_name=..., feature_value=..., ...)
 ```
 
-The wrapped callables seen by `convolve()` can be called like
+The wrapped callables seen by convolve() can be called like
 
 ```python
 wrapped_f(feature_value, ...)
 ```
 
-The first three arguments of `f` are set to the input GraphTensor of
+The first three arguments of f are set to the input GraphTensor of
 the layer and the tag/name pair required by tfgnn.broadcast() and
 tfgnn.pool() to move values between the receiver and the messages that
 are computed inside the convolution. The sole positional argument of
-`wrapped_f()` is passed to `f()`  as `feature_value=`, and any keyword
+wrapped_f() is passed to f()  as feature_value=, and any keyword
 arguments are forwarded.
 
 

 
-`**kwargs`
+**kwargs
 
 
 Forwarded to the base class tf.keras.layers.Layer.
@@ -215,30 +216,31 @@ Forwarded to the base class tf.keras.layers.Layer.
 

 
 
+
  

 
 
Attributes

 
 

 
-`takes_receiver_input`
+takes_receiver_input
 
 
-If `False`, all calls to convolve() will get `receiver_input=None`.
+If False, all calls to convolve() will get receiver_input=None.
 
 

 
-`takes_sender_edge_input`
+takes_sender_edge_input
 
 
-If `False`, all calls to convolve() will get `sender_edge_input=None`.
+If False, all calls to convolve() will get sender_edge_input=None.
 
 

 
-`takes_sender_node_input`
+takes_sender_node_input
 
 
-If `False`, all calls to convolve() will get `sender_node_input=None`.
+If False, all calls to convolve() will get sender_node_input=None.
 
 

 

@@ -247,7 +249,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 
convolve

 
-View
+View
 source
 
 
@@ -273,44 +275,45 @@ from nodes to context). In the end, values have to be pooled from there into a
 Tensor with a leading dimension indexed by receivers, see `pool_to_receiver`.
 
 
+
  

 
 
Args

 
 

 
-`sender_node_input`
+sender_node_input
 
 
-The input Tensor from the sender NodeSet, or `None`.
-If self.takes_sender_node_input is `False`, this arg will be `None`.
-(If it is `True`, that depends on how this layer gets called.)
+The input Tensor from the sender NodeSet, or None.
+If self.takes_sender_node_input is False, this arg will be None.
+(If it is True, that depends on how this layer gets called.)
 See also broadcast_from_sender_node.
 
 

 
-`sender_edge_input`
+sender_edge_input
 
 
-The input Tensor from the sender EdgeSet, or `None`.
-If self.takes_sender_edge_input is `False`, this arg will be `None`.
-(If it is `True`, it depends on how this layer gets called.)
+The input Tensor from the sender EdgeSet, or None.
+If self.takes_sender_edge_input is False, this arg will be None.
+(If it is True, it depends on how this layer gets called.)
 If present, this Tensor is already indexed by the items for which
 messages are computed.
 
 

 
-`receiver_input`
+receiver_input
 
 
 The input Tensor from the receiver NodeSet or Context,
-or None. If self.takes_receiver_input is `False`, this arg will be
-`None`. (If it is `True`, it depends on how this layer gets called.)
+or None. If self.takes_receiver_input is False, this arg will be
+None. (If it is True, it depends on how this layer gets called.)
 See broadcast_from_receiver.
 
 

 
-`broadcast_from_sender_node`
+broadcast_from_sender_node
 
 
 A function that broadcasts a Tensor indexed
@@ -319,25 +322,25 @@ messages are computed.
 
 

 
-`broadcast_from_receiver`
+broadcast_from_receiver
 
 
-Call this as `broadcast_from_receiver(value)`
+Call this as broadcast_from_receiver(value)
 to broadcast a Tensor indexed like receiver_input to a Tensor indexed
 by the items for which messages are computed.
 
 

 
-`pool_to_receiver`
+pool_to_receiver
 
 
-Call this as `pool_to_receiver(value, reduce_type=...)`
+Call this as pool_to_receiver(value, reduce_type=...)
 to pool an item-indexed Tensor to a receiver-indexed tensor, using
 a reduce_type understood by tfgnn.pool(), such as "sum".
 
 

 
-`extra_receiver_ops`
+extra_receiver_ops
 
 
 The extra_receiver_ops passed to init, see there,
@@ -347,10 +350,10 @@ this argument, so subclass implementors not using it can omit it.
 
 

 
-`training`
+training
 
 
-The `training` boolean that was passed to Layer.call(). If true,
+The training boolean that was passed to Layer.call(). If true,
 the result is computed for training rather than inference. For example,
 calls to tf.nn.dropout() are usually conditioned on this flag.
 By contrast, calling another Keras layer (like tf.keras.layers.Dropout)
@@ -360,6 +363,7 @@ does not require forwarding this arg, Keras does that automatically.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SingleInputNextState.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SingleInputNextState.md
index c121434f..5a78ac01 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SingleInputNextState.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/SingleInputNextState.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.SingleInputNextState
 
-[TOC]
-
 
 
-

 
 
Returns

-
-
+
+ View source
+on GitHub 
 
 Replaces a state from a single input.
 
@@ -27,7 +20,11 @@ In a NodeSetUpdate, it replaces the node state with a single edge set input. For
 an EdgeSetUpdate, it replaces the edge_state with the incident node set's input.
 For a ContextUpdate, it replaces the context state with a single node set input.
 
+This layer can be restored from config by `tf.keras.models.load_model()` when
+saved as part of a Keras model using `save_format="tf"`.
+
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/StructuredReadout.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/StructuredReadout.md
index 748d8d83..5e3a0a50 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/StructuredReadout.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/StructuredReadout.md
@@ -1,17 +1,10 @@
 # tfgnn.keras.layers.StructuredReadout
 
-[TOC]
-
 
 
-

 
 
Call returns

-
-
+
+ View source
+on GitHub 
 
 Reads out a feature value from select nodes (or edges) in a graph.
 
@@ -68,16 +61,16 @@ To retrieve a feature unchanged, see
 
 

 
-`key`
+key
 
 
 A string key to select between possibly multiple named readouts
-(such as `"source"` and `"target"` for link prediction). Can be fixed
+(such as "source" and "target" for link prediction). Can be fixed
 in init, or selected for each call.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of the feature to read. If unset (also in call),
@@ -85,15 +78,15 @@ The name of the feature to read. If unset (also in call),
 
 

 
-`readout_node_set`
+readout_node_set
 
 
-A string, defaults to `"_readout"`. This is used as the
+A string, defaults to "_readout". This is used as the
 name for the readout node set and as a name prefix for its edge sets.
 
 

 
-`validate`
+validate
 
 
 Setting this to false disables the validity checks for the
@@ -112,14 +105,14 @@ structurally unchanged GraphTensors.
 
 

 
-`graph`
+graph
 
 
 The scalar GraphTensor to read from.
 
 

 
-`key`
+key
 
 
 Same meaning as for init. Must be passed to init, or to call,
@@ -136,7 +129,7 @@ or to both (with the same value).
 

 
 A tensor of read-out feature values, shaped like a feature of the
-`readout_node_set`.
+readout_node_set.
 
 

-
-
+
+ View source
+on GitHub 
 
 Learns the optimal size constraints for the fixed size batching with retry.
 
@@ -27,20 +20,16 @@ Learns the optimal size constraints for the fixed size batching with retry.
 ) -> Union[tfgnn.SizeConstraints, List[Any]]
 
 
-
-
 
 
 The function estimates the smallest possible size constraints so that a random
-sample of `batch_size` graph tensors meets those constraints with probability
-no less than `success_ratio`. The success ratio is treated as a hard
-constraint, up to sampling error. The constraints can be used for graph tensor
-padding to the fully defined shapes required by XLA.
+sample of `batch_size` graph tensors meets those constraints with probability no
+less than `success_ratio`. The success ratio is treated as a hard constraint, up
+to sampling error. The constraints can be used for graph tensor padding to the
+fully defined shapes required by XLA.
 
 #### Example:
 
-
-
 ```python
 # Learn size constraints for a given dataset of graph tensors and the target
 # batch size(s). The constraints could be learned once and then reused.
@@ -73,35 +62,36 @@ if training:
 
 The learned constraints are intend to be used only with randomized repeated
 dataset. This dataset are first batched using `tf.data.Dataset.batch()`, the
-batches that are too large to fit the learned contraints are filtered using
-tfgnn.satisfies_size_constraints() and then padded
+batches that are too large to fit the learned constraints are filtered using
+tfgnn.satisfies_size_constraints()
+and then padded
 tfgnn.pad_to_total_sizes().
 
 This approach, if applicable, is more efficient compared to padding to the
 maximum possible sizes. It is also simpler and faster compared to the dynamic
-batching, especially for the large batch sizes (>10).  To illustrate the main
+batching, especially for the large batch sizes (>10). To illustrate the main
 point, consider graphs containing only 0 or 1 nodes. A random batch of 1000 of
 those graphs could contain 1000 nodes in the worst case. If this maximum limit
 is used to reseve space for random 1000 graphs, the space of 425 nodes is used
 only in 1:1000_000 cases. It is >40% more efficient to reserve space only for
 575 nodes and resample batches in the rare cases when they do not fit.
 
-
 
+
  

 
 
Args

 
 

 
-`dataset`
+dataset
 
 
 dataset of graph tensors that is intended to be batched.
 
 

 
-`batch_size`
+batch_size
 
 
 the target batch size(s). Could be a single positive integer
@@ -110,7 +100,7 @@ requested value.
 
 

 
-`min_nodes_per_component`
+min_nodes_per_component
 
 
 mapping from a node set name to a minimum number of
@@ -118,7 +108,7 @@ nodes in each graph component. Defaults to 0.
 
 

 
-`success_ratio`
+success_ratio
 
 
 the target probability(s) that a random batch of graph tensor
@@ -126,19 +116,19 @@ satisfies the learned constraints. Could be a single float value between 0
 and 1 or any iterable. For the latter case the result is reported for
 each requested value. NOTE: setting success_ratio to 1 only guarantees
 that all sampled graphs are satisfy the learned constraints. This does not
-in general apply to an arbitrary sample. When `sample_size` tends to
+in general apply to an arbitrary sample. When sample_size tends to
 infinity, the 1 ratio corresponds to the "almost surely satisfies" event.
 
 

 
-`sample_size`
+sample_size
 
 
 the number of the first dataset examples to use for inference.
 
 

 
-`num_thresholds`
+num_thresholds
 
 
 the number of quantiles to use to approximate probability
@@ -148,18 +138,18 @@ distributions.
 

 
 
+
  

 
 
Returns

 

 
-Learned size constraints. If both `batch_size` and `success_ratio` are
-iterables, the result is returned as a nested lists, were `result[b][r]`
-is a size constraints for `batch_size[b]` and `success_ratio[r]`. If any of
-`batch_size` or/and `success_ratio` are scalars the corresponding dimension
+Learned size constraints. If both batch_size and success_ratio are
+iterables, the result is returned as a nested lists, were result[b][r]
+is a size constraints for batch_size[b] and success_ratio[r]. If any of
+batch_size or/and success_ratio are scalars the corresponding dimension
 is squeezed in the output.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/mask_edges.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/mask_edges.md
index 5a15e628..5f2a5691 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/mask_edges.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/mask_edges.md
@@ -1,17 +1,10 @@
 # tfgnn.mask_edges
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Creates a GraphTensor after applying edge_mask over the specified edge-set.
 
@@ -33,35 +26,36 @@ Edge masking doesn't change the node sets or the context node information.
 Not compatible with XLA.
 
 
+
  

 
 
Args

 
 

 
-`graph`
+graph
 
 
 A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 Name of edge-set to apply the boolean mask.
 
 

 
-`boolean_edge_mask`
+boolean_edge_mask
 
 
-A boolean mask with shape `[num_edges]` to mask edges in
-the specified edge-set of the `graph` with rank=1.
+A boolean mask with shape [num_edges] to mask edges in
+the specified edge-set of the graph with rank=1.
 
 

 
-`masked_info_edge_set_name`
+masked_info_edge_set_name
 
 
 Masked out edge-set information will be kept in a
@@ -71,6 +65,7 @@ new edge-set, with name masked_info_edge_set_name.
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md
index 904791d8..1d0979ea 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md
@@ -1,17 +1,10 @@
 # tfgnn.node_degree
 
-[TOC]
-
 
 
-

 
 
Returns

-
-
+
+ View source
+on GitHub 
 
 Returns the degree of each node w.r.t. one side of an edge set.
 
@@ -24,29 +17,29 @@ Returns the degree of each node w.r.t. one side of an edge set.
 
 
 
-
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 The name of the edge set for which degrees are calculated.
 
 

 
-`node_tag`
+node_tag
 
 
 The side of each edge for which the degrees are calculated,
@@ -57,15 +50,16 @@ specified by its tag in the edge set (e.g., t
 

 
 
+
  
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pad_to_total_sizes.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pad_to_total_sizes.md
index 577db1d3..da498e3a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/pad_to_total_sizes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pad_to_total_sizes.md
@@ -1,17 +1,10 @@
 # tfgnn.pad_to_total_sizes
 
-[TOC]
-
 
 
-

 
 
Returns

 

 
-An integer Tensor of shape `[num_nodes]` and dtype equal to `indices_dtype`
-of the GraphTensor. Element `i` contains the number of edges in the given
-edge set that have node index `i` as their endpoint with the given
-`node_tag`. The dimension `num_nodes` is the number of nodes in the
+An integer Tensor of shape [num_nodes] and dtype equal to indices_dtype
+of the GraphTensor. Element i contains the number of edges in the given
+edge set that have node index i as their endpoint with the given
+node_tag. The dimension num_nodes is the number of nodes in the
 respective node set.
 
 

-
-
+
+ View source
+on GitHub 
 
 Pads graph tensor to the total sizes by inserting fake graph components.
 
@@ -25,8 +18,6 @@ Pads graph tensor to the total sizes by inserting fake graph components.
 ) -> Tuple[tfgnn.GraphTensor, tf.Tensor]
 
 
-
-
 
 
 Padding is done by inserting "fake" graph components at the end of the input
@@ -34,58 +25,59 @@ graph tensor until target total sizes are exactly matched. If that is not
 possible (e.g. input already has more nodes than allowed by the constraints)
 function raises `tf.errors.InvalidArgumentError`.
 
-If size_constraints.min_nodes_per_component is specified for a node set,
-the inserted graph components satisfy that constraint (e.g., such that there
-is a node for tf.gather_first_node()). Components in the input graph tensor
-must satisfy that constraint already, or tf.errors.InvalidArgumentError will
-be raised. (This function cannot add padding within existing components.)
+If size_constraints.min_nodes_per_component is specified for a node set, the
+inserted graph components satisfy that constraint (e.g., such that there is a
+node for tf.gather_first_node()). Components in the input graph tensor must
+satisfy that constraint already, or tf.errors.InvalidArgumentError will be
+raised. (This function cannot add padding within existing components.)
 
 Context, node or edge features of the appended fake components are filled with
-user-provided scalar values or with zeros if the latter are not specified.
-Fake edges are created such that each fake node has an approximately uniform
-number of incident edges (this behavior is not guaranteed and may change in
-the future).
+user-provided scalar values or with zeros if the latter are not specified. Fake
+edges are created such that each fake node has an approximately uniform number
+of incident edges (this behavior is not guaranteed and may change in the
+future).
 
 NOTE(b/275338236): This operation is not available in TFLite (last checked for
 TF 2.12).
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 scalar graph tensor (rank=0) to pad.
 
 

 
-`size_constraints`
+size_constraints
 
 
 target total sizes for each graph piece. Must define the
-target number of graph components (`.total_num_components`), target total
-number of items for each node set (`.total_num_nodes[node_set_name]`) and
-likewise for each edge set (`.total_num_edges[edge_set_name]`).
-If `min_nodes_per_component` is set, the inserted graph components satisfy
+target number of graph components (.total_num_components), target total
+number of items for each node set (.total_num_nodes[node_set_name]) and
+likewise for each edge set (.total_num_edges[edge_set_name]).
+If min_nodes_per_component is set, the inserted graph components satisfy
 that constraint and graph components of the input graph tensor are checked
 against this constraint.
 
 

 
-`padding_values`
+padding_values
 
 
 optional mapping from a context, node set or edge set
 feature name to a scalar tensor to use for padding. If no value is
-specified for some feature, its type 'zero' is used (as in `tf.zeros()`).
+specified for some feature, its type 'zero' is used (as in tf.zeros()).
 
 

 
-`validate`
+validate
 
 
 If true, then use assertions to check that the input graph tensor
@@ -97,6 +89,7 @@ each input value.
 

 
 
+
  
@@ -104,35 +97,34 @@ each input value.
 

 
 
Returns

 Tuple of padded graph tensor and padding mask. The mask is a rank-1 dense
 boolean tensor wth size equal to the number of graph compoents is the result
-containing `True` for real graph components and `False` - for fake one used
+containing True for real graph components and False - for fake one used
 for padding.
 
 

 
 

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
 if input parameters are invalid.
 
 

 
-`tf.errors.InvalidArgumentError`
+tf.errors.InvalidArgumentError
 
 
 if input graph tensor could not be padded to
-the `size_constraints` or has less nodes in a component than allowed by
-the `min_nodes_per_component`
+the size_constraints or has less nodes in a component than allowed by
+the min_nodes_per_component
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_example.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_example.md
index 37e77985..b7bcc439 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_example.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_example.md
@@ -1,17 +1,10 @@
 # tfgnn.parse_example
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Parses a batch of serialized Example protos into a single `GraphTensor`.
 
@@ -24,8 +17,6 @@ Parses a batch of serialized Example protos into a single `GraphTensor`.
 ) -> tfgnn.GraphTensor
 
 
-
-
 
 
 We expect `serialized` to be a string tensor batched with `batch_size` many
@@ -38,18 +29,19 @@ function accepts a type spec for a graph tensor and implements an encoding for
 all container tensors, including ragged tensors, from a batched sequence of
 `tf.train.Example` protocol buffer messages.
 
-The encoded examples shapes and features are expected to conform to the
-encoding defined by `get_io_spec()`. The `validate` flag exists to implement
-verifications of this encoding.
+Parsed feature values are converted to their dtype declared in the `spec` using
+`tf.cast()`. Note that `tf.float64` features are converted from a `tf.float32`
+representation in `tf.train.Example`.
 
 
+
  

 
 
Args

 
 

 
-`spec`
+spec
 
 
 A graph tensor type specification of a single serialized graph tensor
@@ -57,16 +49,16 @@ value.
 
 

 
-`serialized`
+serialized
 
 
 A rank-1 dense tensor of strings with serialized Example protos,
-where each example is a graph tensor object with type corresponding `spec`
+where each example is a graph tensor object with type corresponding spec
 type spec.
 
 

 
-`prefix`
+prefix
 
 
 An optional prefix string over all the features. You may use
@@ -74,24 +66,24 @@ this if you are encoding other data in the same protocol buffer.
 
 

 
-`validate`
+validate
 
 
 A boolean indicating whether or not to validate that the input
-values form a valid GraphTensor. Defaults to `True`.
+values form a valid GraphTensor. Defaults to True.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A graph tensor object with `spec.batch(serialized.shape[0])` type spec.
+A graph tensor object with spec.batch(serialized.shape[0]) type spec.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_schema.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_schema.md
index f037f66d..c84a4809 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_schema.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_schema.md
@@ -1,54 +1,46 @@
 # tfgnn.parse_schema
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Parse a schema from text-formatted protos.
 
 
 tfgnn.parse_schema(
     schema_text: Union[bytes, str]
-) -> tfgnn.GraphSchema
+) -> tfgnn.proto.GraphSchema
 

 
 
-
-
 
+
  

 
 
Args

 
 

 
-`schema_text`
+schema_text
 
 
 A string containing a text-formatted protocol buffer rendition
-of a `GraphSchema` message.
+of a GraphSchema message.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A `GraphSchema` instance.
+A GraphSchema instance.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_single_example.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_single_example.md
index 4d98395d..eced5dda 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_single_example.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/parse_single_example.md
@@ -1,17 +1,10 @@
 # tfgnn.parse_single_example
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Parses a single serialized Example proto into a single `GraphTensor`.
 
@@ -24,36 +17,36 @@ Parses a single serialized Example proto into a single `GraphTensor`.
 ) -> tfgnn.GraphTensor
 
 
-
-
 
 
-Like `parse_example()`, but for a single graph tensor.
-See tfgnn.parse_example() for reference.
+Like `parse_example()`, but for a single graph tensor. See
+tfgnn.parse_example() for
+reference.
 
 
+
  

 
 
Args

 
 

 
-`spec`
+spec
 
 
 A graph tensor type specification.
 
 

 
-`serialized`
+serialized
 
 
 A scalar string tensor with a serialized Example proto
-containing a graph tensor object with the `spec` type spec.
+containing a graph tensor object with the spec type spec.
 
 

 
-`prefix`
+prefix
 
 
 An optional prefix string over all the features. You may use
@@ -61,16 +54,17 @@ this if you are encoding other data in the same protocol buffer.
 
 

 
-`validate`
+validate
 
 
 A boolean indicating whether or not to validate that the input
-fields form a valid `GraphTensor`. Defaults to `True`.
+fields form a valid GraphTensor. Defaults to True.
 
 

 

 
 
+
  
@@ -81,4 +75,3 @@ A graph tensor object with a matching type spec.
 

 
 
Returns

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool.md
index 9f2f38ea..c0d229cb 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool.md
@@ -1,17 +1,10 @@
 # tfgnn.pool
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Pools values from edges to nodes, or from nodes or edges to context.
 
@@ -65,20 +58,21 @@ individual results along the innermost axis in the order of appearance.
 support RaggedTensors.
 
 
+
  
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_context.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_context.md
index 497178ab..d0817bf9 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_context.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_context.md
@@ -1,17 +1,10 @@
 # tfgnn.pool_edges_to_context
 
-[TOC]
-
 
 
-

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`to_tag`
+to_tag
 
 
 Values are pooled to context if this is tfgnn.CONTEXT or to the
@@ -86,48 +80,48 @@ incident node on each edge with this tag.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 The name of the edge set from which values are pooled, or
-a non-empty sequence of such names. Unless `to_tag=tfgnn.CONTEXT`,
+a non-empty sequence of such names. Unless to_tag=tfgnn.CONTEXT,
 all named edge sets must have the same incident node set at the given tag.
 
 

 
-`node_set_name`
+node_set_name
 
 
 The name of the node set from which values are pooled,
 or a non-empty sequence of such names. Can only be set with
-`to_tag=tfgnn.CONTEXT`. Exactly one of edge_set_name or node_set_name
+to_tag=tfgnn.CONTEXT. Exactly one of edge_set_name or node_set_name
 must be set.
 
 

 
-`reduce_type`
+reduce_type
 
 
 A string with the name of a pooling operation, or multiple ones
-separated by `|`. See the table above for the known names.
+separated by |. See the table above for the known names.
 
 

 
-`feature_value`
+feature_value
 
 
 A tensor or list of tensors, parallel to the node_set_names
 or edge_set_names, to supply the input values of pooling. Each tensor
-has shape `[num_items, *feature_shape]`, where `num_items` is the number
+has shape [num_items, *feature_shape], where num_items is the number
 of edges in the given edge set or nodes in the given node set, and
-`*feature_shape` is the same across all inputs. The `*feature_shape` may
+*feature_shape is the same across all inputs. The *feature_shape may
 contain ragged dimensions. All the ragged values that are reduced onto
 any one item of the graph must have the same ragged index structure,
 so that a result can be computed from them.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of a feature stored on each graph piece from which
@@ -145,10 +139,10 @@ Exactly one of feature_name or feature_value must be set.
 

 
 A tensor with the result of pooling from the conceptual concatenation of the
-named edge set(s) or node set(s) to the destination selected by `to_tag`.
-Its shape is `[num_items, *feature_shape]`, where `num_items` is the number
-of destination nodes (or graph components if `to_tag=tfgnn.CONTEXT`)
-and `*feature_shape` is as for all the inputs.
+named edge set(s) or node set(s) to the destination selected by to_tag.
+Its shape is [num_items, *feature_shape], where num_items is the number
+of destination nodes (or graph components if to_tag=tfgnn.CONTEXT)
+and *feature_shape is as for all the inputs.
 
 

-
-
+
+ View source
+on GitHub 
 
 Aggregates (pools) edge values to graph context.
 
@@ -49,45 +42,46 @@ regular convolution, we will first broadcast over edges and combine the result
 of that with this function or a pooling over the nodes.)
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 An edge set name.
 
 

 
-`reduce_type`
+reduce_type
 
 
-A pooling operation name, like `"sum"` or `"mean"`, or a
-`|`-separated combination of these; see tfgnn.pool().
+A pooling operation name, like "sum" or "mean", or a
+|-separated combination of these; see tfgnn.pool().
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense edge feature value. Has a shape
-`[num_edges, *feature_shape]`, where `num_edges` is the number of edges in
-the `edge_set_name` edge set and `feature_shape` is the shape of the
+[num_edges, *feature_shape], where num_edges is the number of edges in
+the edge_set_name edge set and feature_shape is the shape of the
 feature value for each edge.
 
 

 
-`feature_name`
+feature_name
 
 
 An edge feature name.
@@ -96,16 +90,16 @@ An edge feature name.
 

 
 
+
  

 
 
Returns

 

 
-A node value pooled to graph context. Has a shape `[num_components,
-*feature_shape]`, where `num_components` is the number of components in a
-graph and `feature_shape` is not affected.
+A node value pooled to graph context. Has a shape [num_components,
+*feature_shape], where num_components is the number of components in a
+graph and feature_shape is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_node.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_node.md
index f8188235..172e53a4 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_node.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_edges_to_node.md
@@ -1,17 +1,10 @@
 # tfgnn.pool_edges_to_node
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Aggregates (pools) edge values to incident nodes.
 
@@ -44,35 +37,36 @@ For a generalization beyond a single edge set, see
 The feature to fetch edge values from is provided either by name (using
 `feature_name`) and found in the graph tensor itself, or provided explicitly
 (using `feature_value`) in which case its shape has to be compatible with the
-shape prefix of the edge set being gathered from. One of `feature_value`
-or `feature_name` must be specified.
+shape prefix of the edge set being gathered from. One of `feature_value` or
+`feature_name` must be specified.
 
 (Note that in most cases the `feature_value` form will be used, because in a
 regular convolution, we will first broadcast over edges and combine the result
 of that with this function.)
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 The name of the edge set from which values are pooled.
 
 

 
-`node_tag`
+node_tag
 
 
 The incident node of each edge at which values are aggregated,
@@ -80,25 +74,25 @@ identified by its tag in the edge set.
 
 

 
-`reduce_type`
+reduce_type
 
 
-A pooling operation name like `"sum"` or `"mean"`, or a
-`|`-separated combination of these; see tfgnn.pool().
+A pooling operation name like "sum" or "mean", or a
+|-separated combination of these; see tfgnn.pool().
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense edge feature value. Has a shape
-`[num_edges, *feature_shape]`, where `num_edges` is the number of edges in
-the `edge_set_name` edge set and `feature_shape` is the shape of the
+[num_edges, *feature_shape], where num_edges is the number of edges in
+the edge_set_name edge set and feature_shape is the shape of the
 feature value for each edge.
 
 

 
-`feature_name`
+feature_name
 
 
 An edge feature name.
@@ -107,16 +101,16 @@ An edge feature name.
 

 
 
+
  

 
 
Returns

 

 
-The edge values pooled to each incident node. Has a shape `[num_nodes,
-*feature_shape]`, where `num_nodes` is the number of nodes in the incident
-node set and `feature_shape` is not affected.
+The edge values pooled to each incident node. Has a shape [num_nodes,
+*feature_shape], where num_nodes is the number of nodes in the incident
+node set and feature_shape is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node.md
new file mode 100644
index 00000000..50f43ccc
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node.md
@@ -0,0 +1,110 @@
+# tfgnn.pool_neighbors_to_node
+
+
+
+
+ View source
+on GitHub 
+
+Aggregates (pools) neighbor node values along one or more edge sets.
+
+
+tfgnn.pool_neighbors_to_node(
+    graph_tensor: GraphTensor,
+    edge_set_name: Union[Sequence[EdgeSetName], EdgeSetName],
+    to_tag: IncidentNodeTag,
+    *,
+    reduce_type: str,
+    from_tag: Optional[IncidentNodeTag] = None,
+    feature_value: Optional[Field] = None,
+    feature_name: Optional[FieldName] = None
+) -> Field
+

+
+
+
+This is a helper function that first broadcasts feature values from the nodes at
+the `from_tag` endpoints (the "neighbors") of each given edge set and then pools
+those values at the `to_tag` endpoints (the "nodes").
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+graph_tensor
+
+A scalar GraphTensor.
+

+edge_set_name
+
+The name of the edge set through which values are pooled, or
+a non-empty sequence of such names. It is required that all edge sets
+connect the same from_tag and to_tag node sets.
+

+to_tag
+
+The incident node of each edge at which values are aggregated,
+identified by its tag in the edge set.
+

+reduce_type
+
+A pooling operation name like "sum" or "mean", or a
+|-separated combination of these; see tfgnn.pool().
+

+from_tag
+
+The incident node of each edge from which values are aggregated.
+Required for hypergraphs. For ordinary graphs, defaults to the opposite of
+to_tag.
+

+feature_value
+
+A ragged or dense neighbor feature value. Has a shape
+[num_sender_nodes, *feature_shape], where num_sender_nodes is the
+number of sender nodes and feature_shape is the shape of the feature
+value for each sender node.
+

+feature_name
+
+An neighbors feature name.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+The sender nodes values pooled to each receiver node. Has a shape
+[num_receiver_nodes, *feature_shape], where num_receiver_nodes is the
+number of receiver nodes and feature_shape is not affected.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node_feature.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node_feature.md
new file mode 100644
index 00000000..d5b086ab
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_neighbors_to_node_feature.md
@@ -0,0 +1,105 @@
+# tfgnn.pool_neighbors_to_node_feature
+
+
+
+
+ View source
+on GitHub 
+
+Aggregates (pools) sender node feature to receiver nodes feature.
+
+
+tfgnn.pool_neighbors_to_node_feature(
+    graph_tensor: GraphTensor,
+    edge_set_name: Union[Sequence[EdgeSetName], EdgeSetName],
+    to_tag: IncidentNodeTag,
+    *,
+    reduce_type: str,
+    feature_name: const.FieldName,
+    to_feature_name: Optional[const.FieldName] = None,
+    from_tag: Optional[IncidentNodeTag] = None
+) -> gt.GraphTensor
+

+
+
+
+Similar to the `pool_neighbors_to_node` but results in the graph tensor with
+updated receiver feature.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Args

+graph_tensor
+
+A scalar GraphTensor.
+

+edge_set_name
+
+The name of the edge set through which values are pooled, or
+a non-empty sequence of such names. It is required that all edge sets
+connect the same from_tag and to_tag node sets.
+

+to_tag
+
+The incident node of each edge at which values are aggregated,
+identified by its tag in the edge set.
+

+reduce_type
+
+A pooling operation name like "sum" or "mean", or a
+|-separated combination of these; see tfgnn.pool().
+

+feature_name
+
+An neighbors feature name to pool values from.
+

+to_feature_name
+
+A receiver feature name to write pooled values to. Defaults
+to the feature_name.
+

+from_tag
+
+The incident node of each edge from which values are aggregated.
+Optional for regular graphs, required for hypergraphs.
+

+
+
+
+ 
+
+
+
+
+
+
+
Returns

+Copy of the input graph tensor with updated receiver_feature_name for the
+receiver node set.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_nodes_to_context.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_nodes_to_context.md
index c8731e07..a33ee821 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_nodes_to_context.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/pool_nodes_to_context.md
@@ -1,17 +1,10 @@
 # tfgnn.pool_nodes_to_context
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Aggregates (pools) node values to graph context.
 
@@ -45,45 +38,46 @@ The feature to fetch node values from is provided either by name (using
 specified.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`node_set_name`
+node_set_name
 
 
 A node set name.
 
 

 
-`reduce_type`
+reduce_type
 
 
-A pooling operation name, like `"sum"` or `"mean"`, or a
-`|`-separated combination of these; see tfgnn.pool().
+A pooling operation name, like "sum" or "mean", or a
+|-separated combination of these; see tfgnn.pool().
 
 

 
-`feature_value`
+feature_value
 
 
 A ragged or dense node feature value. Has a shape
-`[num_nodes, *feature_shape]`, where `num_nodes` is the number of nodes in
-the `node_set_name` node set and `feature_shape` is the shape of the
+[num_nodes, *feature_shape], where num_nodes is the number of nodes in
+the node_set_name node set and feature_shape is the shape of the
 feature value for each node.
 
 

 
-`feature_name`
+feature_name
 
 
 A node feature name.
@@ -92,16 +86,16 @@ A node feature name.
 

 
 
+
  

 
 
Returns

 

 
-Node value pooled to graph context. Has a shape `[num_components,
-*feature_shape]`, where `num_components` is the number of components in a
-graph and `feature_shape` is not affected.
+Node value pooled to graph context. Has a shape [num_components,
+*feature_shape], where num_components is the number of components in a
+graph and feature_shape is not affected.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto.md
new file mode 100644
index 00000000..d79ca235
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto.md
@@ -0,0 +1,68 @@
+# Module: tfgnn.proto
+
+
+
+
+ View source
+on GitHub 
+
+The protocol message (protobuf) types defined by TensorFlow GNN.
+
+This package is automatically included in the top-level tfgnn package:
+
+```
+import tensorflow_gnn as tfgnn
+graph_schema = tfgnn.proto.GraphSchema()
+```
+
+Users are also free to import it separately as
+
+```
+import tensorflow_gnn.proto as tfgnn_proto
+graph_schema = tfgnn_proto.GraphSchema()
+```
+
+...which, together with using its more targeted BUILD dependency, can help to
+shrink the bazel-bin/**/*.runfiles/ directory.
+
+## Classes
+
+[`class BigQuery`](../tfgnn/proto/BigQuery.md): Describes a BigQuery table or
+SQL statement as datasource of a graph piece.
+
+[`class Context`](../tfgnn/proto/Context.md): The schema for the features that
+apply across the entire input graph.
+
+[`class EdgeSet`](../tfgnn/proto/EdgeSet.md): The schema shared by a set of
+edges that connect the same pair of node sets.
+
+[`class Feature`](../tfgnn/proto/Feature.md): The schema entry for a single
+feature.
+
+[`class GraphSchema`](../tfgnn/proto/GraphSchema.md): The top-level container
+for the schema of a graph dataset.
+
+[`class Metadata`](../tfgnn/proto/Metadata.md): Extra information optionally
+provided on a context, node set or edge set.
+
+[`class NodeSet`](../tfgnn/proto/NodeSet.md): The schema shared by a set of
+nodes in the graph.
+
+[`class OriginInfo`](../tfgnn/proto/OriginInfo.md): Metadata about the origin of
+the graph data.
+
+
+
+ 
+
+
+
+
+
+
Other Members
  GraphType   ['UNDEFINED', 'FULL',
+'SUBGRAPH', 'RANDOM_WALKS']
+
+An enumeration of graph types according to the method of creation.
+
+For detailed documentation, see the comments in the graph_schema.proto file.
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery.md
new file mode 100644
index 00000000..76365376
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery.md
@@ -0,0 +1,91 @@
+# tfgnn.proto.BigQuery
+
+
+
+
+ View source
+on GitHub 
+
+Describes a BigQuery table or SQL statement as datasource of a graph piece.
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+read_method
+
+ReadMethod read_method
+

+reshuffle
+
+bool reshuffle
+

+sql
+
+string sql
+

+table_spec
+
+TableSpec table_spec
+

+
+## Child Classes
+
+[`class TableSpec`](../../tfgnn/proto/BigQuery/TableSpec.md)
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Class Variables

+DIRECT_READ
+
+2
+

+EXPORT
+
+1
+

+ReadMethod
+
+['UNSPECIFIED', 'EXPORT', 'DIRECT_READ']
+

+UNSPECIFIED
+
+0
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery/TableSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery/TableSpec.md
new file mode 100644
index 00000000..03f59584
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/BigQuery/TableSpec.md
@@ -0,0 +1,40 @@
+# tfgnn.proto.BigQuery.TableSpec
+
+
+
+
+ View source
+on GitHub 
+
+A ProtocolMessage
+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+dataset
+
+string dataset
+

+project
+
+string project
+

+table
+
+string table
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Context.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Context.md
new file mode 100644
index 00000000..56271b16
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Context.md
@@ -0,0 +1,36 @@
+# tfgnn.proto.Context
+
+
+
+
+ View source
+on GitHub 
+
+The schema for the features that apply across the entire input graph.
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Attributes

+features
+
+repeated FeaturesEntry features
+

+metadata
+
+Metadata metadata
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/EdgeSet.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/EdgeSet.md
new file mode 100644
index 00000000..57413b69
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/EdgeSet.md
@@ -0,0 +1,64 @@
+# tfgnn.proto.EdgeSet
+
+
+
+
+ View source
+on GitHub 
+
+The schema shared by a set of edges that connect the same pair of node sets.
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+context
+
+repeated string context
+

+description
+
+string description
+

+features
+
+repeated FeaturesEntry features
+

+metadata
+
+Metadata metadata
+

+source
+
+string source
+

+target
+
+string target
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Feature.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Feature.md
new file mode 100644
index 00000000..03a652c4
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Feature.md
@@ -0,0 +1,58 @@
+# tfgnn.proto.Feature
+
+
+
+
+ View source
+on GitHub 
+
+The schema entry for a single feature.
+
+
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+description
+
+string description
+

+dtype
+
+DataType dtype
+

+shape
+
+TensorShapeProto shape
+

+source
+
+string source
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/GraphSchema.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/GraphSchema.md
new file mode 100644
index 00000000..741c53ec
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/GraphSchema.md
@@ -0,0 +1,58 @@
+# tfgnn.proto.GraphSchema
+
+
+
+
+ View source
+on GitHub 
+
+The top-level container for the schema of a graph dataset.
+
+
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+context
+
+Context context
+

+edge_sets
+
+repeated EdgeSetsEntry edge_sets
+

+info
+
+OriginInfo info
+

+node_sets
+
+repeated NodeSetsEntry node_sets
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata.md
new file mode 100644
index 00000000..18099619
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata.md
@@ -0,0 +1,54 @@
+# tfgnn.proto.Metadata
+
+
+
+
+ View source
+on GitHub 
+
+Extra information optionally provided on a context, node set or edge set.
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+bigquery
+
+BigQuery bigquery
+

+cardinality
+
+int64 cardinality
+

+extra
+
+repeated KeyValue extra
+

+filename
+
+string filename
+

+
+## Child Classes
+
+[`class KeyValue`](../../tfgnn/proto/Metadata/KeyValue.md)
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata/KeyValue.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata/KeyValue.md
new file mode 100644
index 00000000..0fbd66b3
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/Metadata/KeyValue.md
@@ -0,0 +1,33 @@
+# tfgnn.proto.Metadata.KeyValue
+
+
+
+
+ View source
+on GitHub 
+
+A ProtocolMessage
+
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Attributes

+key
+
+string key
+

+value
+
+string value
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/NodeSet.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/NodeSet.md
new file mode 100644
index 00000000..dd059869
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/NodeSet.md
@@ -0,0 +1,50 @@
+# tfgnn.proto.NodeSet
+
+
+
+
+ View source
+on GitHub 
+
+The schema shared by a set of nodes in the graph.
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
Attributes

+context
+
+repeated string context
+

+description
+
+string description
+

+features
+
+repeated FeaturesEntry features
+

+metadata
+
+Metadata metadata
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/OriginInfo.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/OriginInfo.md
new file mode 100644
index 00000000..404472b2
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/proto/OriginInfo.md
@@ -0,0 +1,36 @@
+# tfgnn.proto.OriginInfo
+
+
+
+
+ View source
+on GitHub 
+
+Metadata about the origin of the graph data.
+
+
+
+For detailed documentation, see the comments in the `graph_schema.proto` file.
+
+
+
+ 
+
+
+
+
+
+
+
+
+
+
+
Attributes

+graph_type
+
+GraphType graph_type
+

+root_set
+
+repeated string root_set
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/random_graph_tensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/random_graph_tensor.md
index de67b023..e9a1ef2b 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/random_graph_tensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/random_graph_tensor.md
@@ -1,32 +1,23 @@
 # tfgnn.random_graph_tensor
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
-Generate a graph tensor from a schema, with random features.
+Generate a graph tensor from a spec, with random features.
 
 
 tfgnn.random_graph_tensor(
     spec: tfgnn.GraphTensorSpec,
     sample_dict: Optional[SampleDict] = None,
     row_lengths_range: Tuple[int, int] = (2, 8),
-    row_splits_dtype: tf.dtypes.DType = tf.int32,
-    validate: bool = True
+    validate: bool = True,
+    num_components_range: Tuple[int, int] = (1, 2)
 ) -> tfgnn.GraphTensor
 

 
-
-
 
 
 NOTE: This function does not (yet?) support the generation of the auxiliary node
@@ -36,20 +27,23 @@ It should not be included in the `spec`, and if needed, should be added
 separately in a later step.
 
 
+
  

 
 
Args

 
 

 
-`spec`
+spec
 
 
-A GraphTensorSpec instance that describes the graph tensor.
+A GraphTensorSpec instance that describes the graph tensor. The result
+random graph tensors are generated for the relaxed number of items, as
+spec.relax(num_components=True, num_nodes=True, num_edges=True).
 
 

 
-`sample_dict`
+sample_dict
 
 
 A dict of (set-type, set-name, field-name) to list-of-values to
@@ -60,40 +54,45 @@ random features are inserted of the right type.
 
 

 
-`row_lengths_range`
+row_lengths_range
 
 
-Minimum and maximum values for each row lengths in a
-ragged range.
+Minimum (included) and maximum (excluded) values for each
+row lengths in a ragged range.
 
 

 
-`row_splits_dtype`
+validate
 
 
-Data type for row splits.
+If true, then use assertions to check that the arguments form a
+valid RaggedTensor. Note: these assertions incur a runtime cost, since
+they must be checked for each tensor value.
 
 

 
-`validate`
+num_components_range
 
 
-If true, then use assertions to check that the arguments form a
-valid RaggedTensor. Note: these assertions incur a runtime cost, since
-they must be checked for each tensor value.
+Minimum (included) and maximum (excluded) values for
+the number of graph components. Overrides the number of components
+from spec.
 
 

 

 
 
+
  

 
 
Returns

 

 
-An instance of a GraphTensor.
+A GraphTensor compatible with the given spec with relaxed number of graph
+items. The size of each node set and edge set is random within
+row_lengths_range. The number of components is random within
+num_components_range.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/read_schema.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/read_schema.md
index 4b159ffd..c60de988 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/read_schema.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/read_schema.md
@@ -1,56 +1,46 @@
 # tfgnn.read_schema
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Read a proto schema from a file with text-formatted contents.
 
 
 tfgnn.read_schema(
     filename: str
-) -> tfgnn.GraphSchema
+) -> tfgnn.proto.GraphSchema
 

 
-
-
 
-
-
 
+
  

 
 
Args

 
 

 
-`filename`
+filename
 
 
 A string, the path to a file containing a text-formatted protocol
-buffer rendition of a `GraphSchema` message.
+buffer rendition of a GraphSchema message.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A `GraphSchema` instance.
+A GraphSchema instance.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/reorder_nodes.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/reorder_nodes.md
index 9c3d574c..00c65716 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/reorder_nodes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/reorder_nodes.md
@@ -1,17 +1,10 @@
 # tfgnn.reorder_nodes
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Reorders nodes within node sets according to indices.
 
@@ -25,84 +18,86 @@ Reorders nodes within node sets according to indices.
 
 
 
-
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`node_indices`
+node_indices
 
 
 A mapping from node sets name to new nodes indices (positions
 within the node set). Each index is an arbitrary permutation of
-`tf.range(num_nodes)`, where `index[i]` is an index of an original node
-to be placed at position `i`.
+tf.range(num_nodes), where index[i] is an index of an original node
+to be placed at position i.
 
 

 
-`validate`
+validate
 
 
-If True, checks that `node_indices` are valid permutations.
+If True, checks that node_indices are valid permutations.
 
 

 

 
 
+
  

 
 
Returns

 

 
-A scalar GraphTensor with randomly shuffled nodes within `node_sets`.
+A scalar GraphTensor with randomly shuffled nodes within node_sets.
 
 

 
 

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If `node_sets` contains non existing node set names.
+If node_sets contains non existing node set names.
 
 

 
-`ValueError`
+ValueError
 
 
-If indices are not `rank=1` `tf.int32` or `tf.int64` tensors.
+If indices are not rank=1 tf.int32 or tf.int64 tensors.
 
 

 
-`InvalidArgumentError`
+InvalidArgumentError
 
 
-if an index shape is not `[num_nodes]`.
+if an index shape is not [num_nodes].
 
 

 
-`InvalidArgumentError`
+InvalidArgumentError
 
 
 if an index is not a permutation of
-`tf.range(num_nodes)`. Only if validate is set to True.
+tf.range(num_nodes). Only if validate is set to True.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/reverse_tag.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/reverse_tag.md
index 5bcff3f4..df2f0b8f 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/reverse_tag.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/reverse_tag.md
@@ -1,17 +1,10 @@
 # tfgnn.reverse_tag
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Flips tfgnn.SOURCE to tfgnn.TARGET and vice versa.
 
@@ -21,6 +14,4 @@ Flips tfgnn.SOURCE to tfgnn.TARGET and vice versa.
 )
 
 
-
-
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler.md
index a97e80ca..14090d56 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler.md
@@ -1,17 +1,10 @@
 # Module: tfgnn.sampler
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Public interface for GNN Sampler.
 
@@ -28,3 +21,19 @@ builder pattern that eases creation of `tfgnn.SamplingSpec`.
 
 [`make_sampling_spec_tree(...)`](../tfgnn/sampler/make_sampling_spec_tree.md):
 Automatically creates `SamplingSpec` by starting from seed node set.
+
+
+
+ 
+
+
+
+
+
+
+
+
Other Members

+SamplingStrategy
+
+['TOP_K', 'RANDOM_UNIFORM', 'RANDOM_WEIGHTED']
+

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingOp.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingOp.md
index 233dd4c5..f25ff6c1 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingOp.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingOp.md
@@ -1,17 +1,10 @@
 # tfgnn.sampler.SamplingOp
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A ProtocolMessage
 
@@ -24,38 +17,38 @@ A ProtocolMessage
 
 

 
-`edge_set_name`
+edge_set_name
 
 
-`string edge_set_name`
+string edge_set_name
 
 

 
-`input_op_names`
+input_op_names
 
 
-`repeated string input_op_names`
+repeated string input_op_names
 
 

 
-`op_name`
+op_name
 
 
-`string op_name`
+string op_name
 
 

 
-`sample_size`
+sample_size
 
 
-`int32 sample_size`
+int32 sample_size
 
 

 
-`strategy`
+strategy
 
 
-`SamplingStrategy strategy`
+SamplingStrategy strategy
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpec.md
index ad0ca821..8d83211c 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpec.md
@@ -1,17 +1,10 @@
 # tfgnn.sampler.SamplingSpec
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 A ProtocolMessage
 
@@ -24,17 +17,24 @@ A ProtocolMessage
 
 

 
-`sampling_ops`
+sampling_ops
+
+repeated SamplingOp sampling_ops
+

+seed_op
 
 
-`repeated SamplingOp sampling_ops`
+SeedOp seed_op
 
 

 
-`seed_op`
+symmetric_link_seed_op
 
 
-`SeedOp seed_op`
+SymmetricLinkSeedOp symmetric_link_seed_op
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpecBuilder.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpecBuilder.md
index 11033cc2..f2926a0a 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpecBuilder.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/sampler/SamplingSpecBuilder.md
@@ -1,24 +1,17 @@
 # tfgnn.sampler.SamplingSpecBuilder
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Mimics builder pattern that eases creation of `tfgnn.SamplingSpec`.
 
 
 tfgnn.sampler.SamplingSpecBuilder(
-    graph_schema: tfgnn.GraphSchema,
-    default_strategy: SamplingStrategy = SamplingStrategy.TOP_K
+    graph_schema: tfgnn.proto.GraphSchema,
+    default_strategy: tfgnn.sampler.SamplingStrategy = SamplingStrategy.TOP_K
 )
 

 
@@ -109,12 +102,12 @@ Initializes sampling by seeding on node with `node_set_name`.
 
 

 
-`node_set_name`
+node_set_name
 
 
-Becomes the `node_set_name` of built `spec.sampling_op`. If
-not given, the graph schema must be homogeneous (with one `node_set`).
-If given, it must correspond to some node set name in `graph_schema`
+Becomes the node_set_name of built spec.sampling_op. If
+not given, the graph schema must be homogeneous (with one node_set).
+If given, it must correspond to some node set name in graph_schema
 given to constructor.
 
 

 
 Object which support builder pattern, upon which, you may repeatedly call
-`.sample()`, per header comments.
+.sample(), per header comments.
 
 

-
-
+
+ View source
+on GitHub 
 
 Automatically creates `SamplingSpec` by starting from seed node set.
 
 
 tfgnn.sampler.make_sampling_spec_tree(
-    graph_schema: tfgnn.GraphSchema,
+    graph_schema: tfgnn.proto.GraphSchema,
     seed_nodeset: NodeSetName,
     *,
     sample_sizes: List[int],
@@ -40,14 +33,14 @@ created `SamplingSpec` instructs sampling up to `sample_sizes[1]` edges for
 
 

 
-`graph_schema`
+graph_schema
 
 
 contains node-sets & edge-sets.
 
 

 
-`seed_nodeset`
+seed_nodeset
 
 
 name of node-set that the sampler will be instructed to use as
@@ -55,17 +48,17 @@ seed nodes.
 
 

 
-`sample_sizes`
+sample_sizes
 
 
-list of number of nodes to sample. E.g. if `sample_sizes` are
-`[5, 2, 2]`, then for every sampled node, up-to `5` of its neighbors will
-be sampled, and for each, up to `2` of its neighbors will be sampled, etc,
-totalling sampled nodes up to `5 * 2 * 2 = 20` for each seed node.
+list of number of nodes to sample. E.g. if sample_sizes are
+[5, 2, 2], then for every sampled node, up-to 5 of its neighbors will
+be sampled, and for each, up to 2 of its neighbors will be sampled, etc,
+totalling sampled nodes up to 5 * 2 * 2 = 20 for each seed node.
 
 

 
-`sampling_strategy`
+sampling_strategy
 
 
 one of the supported sampling strategies, the same for
@@ -81,8 +74,8 @@ each depth.
 
Returns

 

 
-`SamplingSpec` that instructs the sampler to sample according to the
-`sampling_strategy` and `sample_sizes`.
+SamplingSpec that instructs the sampler to sample according to the
+sampling_strategy and sample_sizes.
 
 

-
-
+
+ View source
+on GitHub 
 
 Returns whether the input `graph_tensor` satisfies `total_sizes`.
 
@@ -30,26 +23,23 @@ Returns whether the input `graph_tensor` satisfies `total_sizes`.
 ) -> tf.Tensor
 

 
-
-
 
-
-
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 a graph tensor to check against target total sizes.
 
 

 
-`total_sizes`
+total_sizes
 
 
 target total sizes for each graph piece.
@@ -58,15 +48,15 @@ target total sizes for each graph piece.
 

 
 
+
  

 
 
Returns

 

 
-A scalar boolean tensor equal to `True` if the `graph_tensor` statisifies
-`total_sizes`, and `False` if not.
+A scalar boolean tensor equal to True if the graph_tensor statisifies
+total_sizes, and False if not.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_features_globally.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_features_globally.md
index eef56412..c474ad61 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_features_globally.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_features_globally.md
@@ -1,17 +1,10 @@
 # tfgnn.shuffle_features_globally
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Shuffles context, node set and edge set features of a scalar GraphTensor.
 
@@ -34,14 +27,14 @@ TF 2.12).
 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`seed`
+seed
 
 
 A seed for random uniform shuffle.
@@ -56,15 +49,15 @@ A seed for random uniform shuffle.
 
Returns

 

 
-A scalar GraphTensor `result` with the same graph structure as the input,
+A scalar GraphTensor result with the same graph structure as the input,
 but randomly shuffled feature tensors. More precisely, the result satisfies
-`result.node_sets[ns][ft][i] = graph_tensor.node_sets[ns][ft][sigma(i)]`
-for all node set names `ns` (including auxiliary node sets), all feature
-names `ft` and all indices `i` in `range(n)`, where `n` is the total_size
-of the node set and `sigma` is a permutation of `range(n)`.
+result.node_sets[ns][ft][i] = graph_tensor.node_sets[ns][ft][sigma(i)]
+for all node set names ns (including auxiliary node sets), all feature
+names ft and all indices i in range(n), where n is the total_size
+of the node set and sigma is a permutation of range(n).
 Moreover, the result satisfies the the analogous equations for all features
 of all edge sets (including auxiliary edge sets) and the context.
-The permutation `sigma` is drawn uniformly at random, independently for
+The permutation sigma is drawn uniformly at random, independently for
 each graph piece and each feature(!). That is, separate features are
 permuted differently, and features on any one item (edge, node, component)
 can form combinations not seen on an input item.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_nodes.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_nodes.md
index a58c9d9f..f62162d2 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_nodes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_nodes.md
@@ -1,17 +1,10 @@
 # tfgnn.shuffle_nodes
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Randomly reorders nodes of given node sets, within each graph component.
 
@@ -40,28 +33,29 @@ NOTE(b/277938756): This operation is not available in TFLite (last checked for
 TF 2.12).
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`node_sets`
+node_sets
 
 
 An optional collection of node sets names to shuffle. If None,
-all node sets are shuffled.  Should not overlap with `shuffle_indices`.
+all node sets are shuffled.  Should not overlap with shuffle_indices.
 
 

 
-`seed`
+seed
 
 
 Optionally, a fixed seed for random uniform shuffle.
@@ -70,28 +64,30 @@ Optionally, a fixed seed for random uniform shuffle.
 

 
 
+
  

 
 
Returns

 

 
-A scalar GraphTensor with randomly shuffled nodes within `node_sets`.
+A scalar GraphTensor with randomly shuffled nodes within node_sets.
 
 

 
 

 
 
+
  

 
 
Raises

 
 

 
-`ValueError`
+ValueError
 
 
-If `node_sets` containes non existing node set names.
+If node_sets containes non existing node set names.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax.md
index 5b74662f..caccb2ac 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax.md
@@ -1,17 +1,10 @@
 # tfgnn.softmax
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Computes softmax over a many-to-one relationship in a GraphTensor.
 
@@ -31,28 +24,29 @@ Computes softmax over a many-to-one relationship in a GraphTensor.
 
 This function can be used to compute a softmax normalization...
 
-  * of edge values, across the edges with a common incident node at `per_tag`
+*   of edge values, across the edges with a common incident node at `per_tag`
     (e.g., SOURCE or TARGET);
-  * of node values, across all the nodes in the same graph component;
-  * of edge values, across all the edges in the same graph component.
+*   of node values, across all the nodes in the same graph component;
+*   of edge values, across all the edges in the same graph component.
 
 For non-scalar values, the softmax function is applied element-wise.
 
 
+
  

 
 
Args

 
 

 
-`graph_tensor`
+graph_tensor
 
 
 A scalar GraphTensor.
 
 

 
-`per_tag`
+per_tag
 
 
 tfgnn.CONTEXT for normalization per graph component, or an incident
@@ -61,37 +55,37 @@ common incident node.
 
 

 
-`edge_set_name`
+edge_set_name
 
 
 The name of the edge set on which values are normalized,
-or a non-empty sequence of such names. Unless `from_tag=tfgnn.CONTEXT`,
+or a non-empty sequence of such names. Unless from_tag=tfgnn.CONTEXT,
 all named edge sets must have the same incident node set at the given tag.
 
 

 
-`node_set_name`
+node_set_name
 
 
 The name of the node set on which values are normalized,
 or a non-empty sequence of such names. Can only be passed together with
-`from_tag=tfgnn.CONTEXT`. Exactly one of edge_set_name or node_set_name
+from_tag=tfgnn.CONTEXT. Exactly one of edge_set_name or node_set_name
 must be set.
 
 

 
-`feature_value`
+feature_value
 
 
 A tensor or list of tensors, parallel to the node_set_names
 or edge_set_names, to supply the input values of softmax. Each tensor
-has shape `[num_items, *feature_shape]`, where `num_items` is the number
+has shape [num_items, *feature_shape], where num_items is the number
 of edges in the given edge set or nodes in the given node set, and
-`*feature_shape` is the same across all inputs.
+*feature_shape is the same across all inputs.
 
 

 
-`feature_name`
+feature_name
 
 
 The name of a feature stored on each graph piece from which
@@ -114,4 +108,3 @@ the tensors and the length of the list do not change from the input.
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax_edges_per_node.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax_edges_per_node.md
index daeaa2b0..75e79399 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax_edges_per_node.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/softmax_edges_per_node.md
@@ -1,17 +1,10 @@
 # tfgnn.softmax_edges_per_node
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Returns softmax() of edge values per common `node_tag` node.
 
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout.md
index 861859a6..337242b0 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/structured_readout.md
@@ -1,17 +1,10 @@
 # tfgnn.structured_readout
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Reads out a feature value from select nodes (or edges) in a graph.
 
@@ -89,7 +82,7 @@ as usual with `GraphTensor.node_sets["_readout"]["ft"]`.
 
 

 
-`graph`
+graph
 
 
 A scalar GraphTensor with a readout structure composed of auxiliary
@@ -97,36 +90,36 @@ graph pieces as described above.
 
 

 
-`key`
+key
 
 
 A string key to select between possibly multiple named readouts
-(such as `"source"` and `"target"` for link prediction).
+(such as "source" and "target" for link prediction).
 
 

 
-`feature_name`
+feature_name
 
 
 The name of a feature that is present on the node set(s)
 (or edge set(s)) referenced by the auxiliary edge sets. The feature
-must have shape `[num_items, *feature_dims]` with the same `feature_dims`
+must have shape [num_items, *feature_dims] with the same feature_dims
 on all graph pieces, and the same dtype.
 
 

 
-`readout_node_set`
+readout_node_set
 
 
 The name for the readout node set and the name prefix for
-its edge sets. Permissible values are `"_readout"` (the default) and
-`f"_readout:{tag}"` where `tag` matches `[a-zA-Z0-9_]+`.
+its edge sets. Permissible values are "_readout" (the default) and
+f"_readout:{tag}" where tag matches [a-zA-Z0-9_]+.
 Setting this to a different value allows to select between multiple
 independent readout structures in the same graph.
 
 

 
-`validate`
+validate
 
 
 Setting this to false disables the validity checks for the
@@ -144,7 +137,7 @@ structurally unchanged GraphTensors.
 
Returns

 

 
-A tensor of shape `[readout_size, *feature_dims]` with the read-out feature
+A tensor of shape [readout_size, *feature_dims] with the read-out feature
 values.
 
 

-
-
+
+ View source
+on GitHub 
 
 Reads out a feature value from select nodes (or edges) in a graph.
 
@@ -45,7 +38,7 @@ modified `GraphTensor` in which the readout result is stored as a feature on the
 
 

 
-`graph`
+graph
 
 
 A scalar GraphTensor with the auxiliary graph pieces required by
@@ -53,15 +46,15 @@ A scalar GraphTensor with the auxiliary graph pieces required by
 
 

 
-`key`
+key
 
 
 A string key to select between possibly multiple named readouts
-(such as `"source"` and `"target"` for link prediction).
+(such as "source" and "target" for link prediction).
 
 

 
-`feature_name`
+feature_name
 
 
 The name of a feature to read out from, as with
@@ -69,43 +62,43 @@ The name of a feature to read out from, as with
 
 

 
-`new_feature_name`
+new_feature_name
 
 
-The name of the feature to add to `readout_node_set`
-for storing the readout result. If unset, defaults to `feature_name`.
-It is an error if the added feature already exists on `readout_node_set`
-in the input `graph`, unless `overwrite=True` is set.
+The name of the feature to add to readout_node_set
+for storing the readout result. If unset, defaults to feature_name.
+It is an error if the added feature already exists on readout_node_set
+in the input graph, unless overwrite=True is set.
 
 

 
-`remove_input_feature`
+remove_input_feature
 
 
-If set, the given `feature_name` is removed from the
+If set, the given feature_name is removed from the
 node (or edge) set(s) that supply the input to
 tfgnn.structured_readout().
 
 

 
-`overwrite`
+overwrite
 
 
 If set, allows overwriting a potentially already existing
-feature `graph.node_sets[readout_node_set][new_feature_name]`.
+feature graph.node_sets[readout_node_set][new_feature_name].
 
 

 
-`readout_node_set`
+readout_node_set
 
 
-A string, defaults to `"_readout"`. This is used as the
+A string, defaults to "_readout". This is used as the
 name for the readout node set and as a name prefix for its edge sets.
 See tfgnn.structured_readout() for more.
 
 

 
-`validate`
+validate
 
 
 Setting this to false disables the validity checks of
@@ -123,9 +116,9 @@ earlier on structurally unchanged GraphTensors.
 
Returns

 

 
-A `GraphTensor` like `graph`, with the readout result stored as
-`.node_sets[readout_node_set][new_feature_name]` and possibly the
-readout inputs removed (see `remove_input_feature`).
+A GraphTensor like graph, with the readout result stored as
+.node_sets[readout_node_set][new_feature_name] and possibly the
+readout inputs removed (see remove_input_feature).
 
 

-
-
+
+ View source
+on GitHub 
 
 Checks `graph` supports `structured_readout()` from `required_keys`.
 
@@ -45,14 +38,14 @@ all keys provided in the graph structure are checked.
 
 

 
-`graph`
+graph
 
 
 The graph tensor to check.
 
 

 
-`required_keys`
+required_keys
 
 
 Can be set to a list of readout keys to check. If unset,
@@ -60,7 +53,7 @@ checks all keys provided by the graph.
 
 

 
-`readout_node_set`
+readout_node_set
 
 
 Optionally, a non-default name for use as
@@ -76,7 +69,7 @@ Optionally, a non-default name for use as
 
Returns

 

 
-The input GraphTensor, unchanged. This helps to put `tf.debugging.assert*`
+The input GraphTensor, unchanged. This helps to put tf.debugging.assert*
 ops from this function into a dependency chain.
 
 

 
-`ValueError`
+ValueError
 
 
 if the auxiliary graph pieces for readout are malformed in the
@@ -99,14 +92,14 @@ if the auxiliary graph pieces for readout are malformed in the
 
 

 
-`KeyError`
+KeyError
 
 
-if any of the `required_keys` is missing.
+if any of the required_keys is missing.
 
 

 
-`tf.errors.InvalidArgumentError`
+tf.errors.InvalidArgumentError
 
 
 If values in the GraphTensor, notably
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_spec_for_readout.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_spec_for_readout.md
index cb656e22..c3429762 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_spec_for_readout.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_graph_tensor_spec_for_readout.md
@@ -1,17 +1,10 @@
 # tfgnn.validate_graph_tensor_spec_for_readout
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Checks `graph_spec` supports `structured_readout()` from `required_keys`.
 
@@ -53,7 +46,7 @@ input.
 
 

 
-`graph_spec`
+graph_spec
 
 
 The graph tensor spec to check. Must be scalar, that is, have
@@ -61,7 +54,7 @@ shape [].
 
 

 
-`required_keys`
+required_keys
 
 
 Can be set to a list of readout keys that are required to be
@@ -69,7 +62,7 @@ provided by the spec.
 
 

 
-`readout_node_set`
+readout_node_set
 
 
 The name of the auxiliary node set for readout, which is
@@ -87,17 +80,17 @@ auxiliary edge sets connected to it.
 
 

 
-`ValueError`
+ValueError
 
 
 if the auxiliary graph pieces for readout are malformed.
 
 

 
-`KeyError`
+KeyError
 
 
-if any of the `required_keys` is missing.
+if any of the required_keys is missing.
 
 

 

diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_schema.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_schema.md
index ef13c26a..5fa193bf 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_schema.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/validate_schema.md
@@ -1,23 +1,16 @@
 # tfgnn.validate_schema
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Validates the correctness of a graph schema instance.
 
 
 tfgnn.validate_schema(
-    schema: tfgnn.GraphSchema,
+    schema: tfgnn.proto.GraphSchema,
     readout_node_sets: Optional[Sequence[const.NodeSetName]] = None
 ) -> List[Exception]
 

@@ -25,36 +18,38 @@ Validates the correctness of a graph schema instance.
 
 
 `GraphSchema` configuration messages are created by users in order to describe
-the topology of a graph. This function checks various aspects of the schema
-for correctness, e.g. prevents usage of reserved feature names, ensures given
-shapes are fully-defined, ensures set name references are found, etc.
+the topology of a graph. This function checks various aspects of the schema for
+correctness, e.g. prevents usage of reserved feature names, ensures given shapes
+are fully-defined, ensures set name references are found, etc.
 
 
+
  

 
 
Args

 
 

 
-`schema`
+schema
 
 
 An instance of the graph schema.
 
 

 
-`readout_node_sets`
+readout_node_sets
 
 
 By default, this function checks the "_readout" node set,
 if present, if it meets the requirements of tfgnn.structured_readout().
 That's sufficient for most cases. Optionally, you can pass a list of
-`readout_node_set` names to (a) require their presence and (b) check them.
+readout_node_set names to (a) require their presence and (b) check them.
 
 

 

 
 
+
  
@@ -67,16 +62,15 @@ Render those to your favorite stream (or ignore).
 
 

 
 
Returns

 
-
-
 
+
  

 
 
Raises

 
 

 
-`ValidationError`
+ValidationError
 
 
 If a validation check fails.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/write_example.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/write_example.md
index e2279ed5..321be541 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/write_example.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/write_example.md
@@ -1,17 +1,10 @@
 # tfgnn.write_example
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Encode an eager `GraphTensor` to a tf.train.Example proto.
 
@@ -22,8 +15,6 @@ Encode an eager `GraphTensor` to a tf.train.Example proto.
 ) -> tf.train.Example
 
 
-
-
 
 
 This routine can be used to create a stream of training data for GNNs from a
@@ -39,21 +30,31 @@ graph schema. The graph tensors materialized in this way will be parseable by
 (using the spec deserialized from the schema) and have the same contents (up to
 the choice of indices_dtype).
 
+All features stored on the `graph` must have dtypes that are supported by the
+graph schema. For the following dtypes, special caveats apply to their
+representation in `tf.train.Example`:
+
+*   `tf.float64` features are serialized as `tf.float32`, which truncates them
+    (perhaps fatally to 0 or +/- inf, if exceeding the exponent range).
+*   `tf.uint64` features are serialized as `tf.int64` values with the same bit
+    pattern. Deserializing from `tf.train.Example` recovers the original value.
+
 
+
  

 
 
Args

 
 

 
-`graph`
+graph
 
 
-An eager instance of `GraphTensor` to write out.
+An eager instance of GraphTensor to write out.
 
 

 
-`prefix`
+prefix
 
 
 An optional prefix string over all the features. You may use
@@ -63,15 +64,14 @@ this if you are encoding other data in the same protocol buffer.
 

 
 
+
  

 
 
Returns

 

 
-A reference to `result`, if provided, or a to a freshly created instance
-of `tf.train.Example`.
+A tf.train.Example with the serialized graph.
 
 

 
 

-
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/write_schema.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/write_schema.md
index 6d034f9f..a4c7c325 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/write_schema.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/write_schema.md
@@ -1,51 +1,41 @@
 # tfgnn.write_schema
 
-[TOC]
-
 
 
-
-
-
+
+ View source
+on GitHub 
 
 Write a `GraphSchema` to a text-formatted proto file.
 
 
 tfgnn.write_schema(
-    schema: tfgnn.GraphSchema,
+    schema: tfgnn.proto.GraphSchema,
     filename: str
 )
 

 
-
-
 
-
-
 
+
  

 
 
Args

 
 

 
-`schema`
+schema
 
 
-A `GraphSchema` instance to write out.
+A GraphSchema instance to write out.
 
 

 
-`filename`
+filename
 
 
 A string, the path to a file to render a text-formatted rendition
-of the `GraphSchema` message to.
+of the GraphSchema message to.
 
 

 

diff --git a/tensorflow_gnn/models/multi_head_attention/layers.py b/tensorflow_gnn/models/multi_head_attention/layers.py
index 9693c8a8..8f1d9089 100644
--- a/tensorflow_gnn/models/multi_head_attention/layers.py
+++ b/tensorflow_gnn/models/multi_head_attention/layers.py
@@ -84,6 +84,8 @@ class MultiHeadAttentionConv(tfgnn.keras.layers.AnyToAnyConvolutionBase):
   Example: Transformer-style attention on neighbors along incoming edges
   whose result is concatenated with the old node state and passed through
   a Dense layer to compute the new node state.
+
+
   ```
   dense = tf.keras.layers.Dense
   graph = tfgnn.keras.layers.GraphUpdate(
@@ -97,22 +99,22 @@ class MultiHeadAttentionConv(tfgnn.keras.layers.AnyToAnyConvolutionBase):
   For now, there is a variant that modifies the inputs transformation part and
   could potentially be beneficial:
 
-      1. (transform_keys is False) Instead of projecting both queries and
-        keys when computing attention weights, we only project the queries
-        because the two linear projections can be collapsed to a single
-        projection:
+    1. (transform_keys is False) Instead of projecting both queries and
+       keys when computing attention weights, we only project the queries
+       because the two linear projections can be collapsed to a single
+       projection:
 
-          $$ (Q_v W_Q^k)(K_u W_K^k)^T
-            = Q_v (W_Q^k {W_K^k}^T) K_u^T
-            = Q_v W_{QK}^k K_u^T $$
+         $$ (Q_v W_Q^k)(K_u W_K^k)^T
+           = Q_v (W_Q^k {W_K^k}^T) K_u^T
+           = Q_v W_{QK}^k K_u^T $$
 
-        where $d$ is the key width. (Following "Attention is all you need",
-        this scaling is meant to achieve unit variance of the results, assuming
-        that $Q_v W_{QK}^k$ has unit variance due to the initialization of
-        $Q_v W_{QK}^k$.)
+       where $d$ is the key width. (Following "Attention is all you need",
+       this scaling is meant to achieve unit variance of the results, assuming
+       that $Q_v W_{QK}^k$ has unit variance due to the initialization of
+       $Q_v W_{QK}^k$.)
 
-        NOTE: The single projection matrix behaves differently in
-        gradient-descent training than the product of two matrices.
+       NOTE: The single projection matrix behaves differently in
+       gradient-descent training than the product of two matrices.
 
   This layer can be restored from config by `tf.keras.models.load_model()`
   when saved as part of a Keras model using `save_format="tf"`.