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 cc01470b..1b329613 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
@@ -20,7 +20,7 @@ Returns a GraphUpdate layer with a Graph Attention Network V2 (GATv2).
     *,
     num_heads: int,
     per_head_channels: int,
-    receiver_tag: tfgnn.IncidentNodeOrContextTag,
+    receiver_tag: tfgnn.IncidentNodeTag,
     feature_name: str = tfgnn.HIDDEN_STATE,
     heads_merge_type: str = 'concat',
     name: str = 'gat_v2',
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 d3cf1a08..408b4335 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
@@ -22,7 +22,7 @@ Returns a GraphUpdate layer for message passing with GATv2 pooling.
     message_dim: int,
     num_heads: int,
     heads_merge_type: str = 'concat',
-    receiver_tag: tfgnn.IncidentNodeOrContextTag,
+    receiver_tag: tfgnn.IncidentNodeTag,
     node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None,
     edge_feature: Optional[tfgnn.FieldName] = None,
     l2_regularization: float = 0.0,
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 a738e0b5..2df82a19 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNConv.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNConv.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/gcn/gcn_conv.py#L28-L191">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/gcn/gcn_conv.py#L28-L218">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -27,6 +27,7 @@ Implements the Graph Convolutional Network by Kipf&Welling (2016).
     kernel_initializer: bool = None,
     node_feature: Optional[str] = tfgnn.HIDDEN_STATE,
     kernel_regularizer: Optional[_RegularizerType] = None,
+    edge_weight_feature_name: Optional[tfgnn.FieldName] = None,
     **kwargs
 )
 </code></pre>
@@ -84,7 +85,8 @@ with Keras and other implementations.
 </td>
 <td>
 Whether to compute the result as if a loop from each node
-to itself had been added to the edge set.
+to itself had been added to the edge set. The self-loop edges are added
+with an edge weight of one.
 </td>
 </tr><tr>
 <td>
@@ -109,6 +111,15 @@ Name of the node feature to transform.
 </td>
 </tr><tr>
 <td>
+`edge_weight_feature_name`<a id="edge_weight_feature_name"></a>
+</td>
+<td>
+Can be set to the name of a feature on the edge
+set that supplies a scalar weight for each edge. The GCN computation uses
+it as the edge's entry in the adjacency matrix, instead of the default 1.
+</td>
+</tr><tr>
+<td>
 `**kwargs`<a id="**kwargs"></a>
 </td>
 <td>
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 aa152834..9da70620 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNHomGraphUpdate.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/gcn/GCNHomGraphUpdate.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/gcn/gcn_conv.py#L194-L251">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/gcn/gcn_conv.py#L221-L278">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
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 490dc0f8..5f71c0e3 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L253-L469">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L257-L473">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
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 bdede4a7..eb9d6230 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L25-L122">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L25-L124">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -137,7 +137,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 <h3 id="convolve"><code>convolve</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L109-L122">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L109-L124">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -149,6 +149,7 @@ source</a>
     broadcast_from_sender_node: Callable[[tf.Tensor], tf.Tensor],
     broadcast_from_receiver: Callable[[tf.Tensor], tf.Tensor],
     pool_to_receiver: Callable[..., tf.Tensor],
+    extra_receiver_ops: Any = None,
     training: bool
 ) -> tf.Tensor
 </code></pre>
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 9831a117..feafc94c 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L472-L585">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L476-L589">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
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 5a11f269..ab3307e0 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L588-L739">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L592-L743">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
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 cefbdf9b..8b019f2a 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L125-L250">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L127-L254">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -173,7 +173,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 <h3 id="convolve"><code>convolve</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L235-L250">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/graph_sage/layers.py#L237-L254">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -185,6 +185,7 @@ source</a>
     broadcast_from_sender_node: Callable[[tf.Tensor], tf.Tensor],
     broadcast_from_receiver: Callable[[tf.Tensor], tf.Tensor],
     pool_to_receiver: Callable[..., tf.Tensor],
+    extra_receiver_ops: Any = None,
     training: bool
 ) -> tf.Tensor
 </code></pre>
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 957bedb9..47ae18eb 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
@@ -37,3 +37,9 @@ Returns a GraphUpdate layer with a transformer-style multihead attention.
 
 [`MultiHeadAttentionMPNNGraphUpdate(...)`](./multi_head_attention/MultiHeadAttentionMPNNGraphUpdate.md):
 Returns a GraphUpdate layer for message passing with MultiHeadAttention pooling.
+
+[`graph_update_from_config_dict(...)`](./multi_head_attention/graph_update_from_config_dict.md):
+Returns a MultiHeadAttentionMPNNGraphUpdate initialized from `cfg`.
+
+[`graph_update_get_config_dict(...)`](./multi_head_attention/graph_update_get_config_dict.md):
+Returns ConfigDict for graph_update_from_config_dict() with defaults.
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 a3b22963..c90d835d 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L23-L518">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L24-L558">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -31,7 +31,7 @@ Transformer-style (dot-product) multi-head attention on GNNs.
     kernel_initializer: Union[None, str, tf.keras.initializers.Initializer] = None,
     kernel_regularizer: Union[None, str, tf.keras.regularizers.Regularizer] = None,
     transform_keys: bool = True,
-    score_scaling: bool = True,
+    score_scaling: Literal['none', 'rsqrt_dim', 'trainable_sigmoid'] = &#x27;rsqrt_dim&#x27;,
     transform_values_after_pooling: bool = False,
     **kwargs
 )
@@ -85,13 +85,13 @@ Note that in the context of graph, only nodes with edges connected are attended
 to each other, which means we do NOT compute $N^2$ pairs of scores as the
 original Transformer-style Attention.
 
-Users are able to remove the scaling of attention scores (score_scaling=False)
-or add an activation on the transformed query (controled by
-`attention_activation`). However, we recommend to remove the scaling when using
-an `attention_activation` since activating both of them may lead to degrated
-accuracy. One can also customize the transformation kernels with different
-intializers, regularizers as well as the use of bias terms, using the other
-arguments.
+Users are able to remove the scaling of attention scores
+(`score_scaling="none"`) or add an activation on the transformed query
+(controlled by `attention_activation`). However, we recommend to remove the
+scaling when using an `attention_activation` since activating both of them may
+lead to degraded accuracy. One can also customize the transformation kernels
+with different initializers, regularizers as well as the use of bias terms,
+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
@@ -125,7 +125,6 @@ could potentially be beneficial:
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Init args</h2></th></tr>
@@ -256,9 +255,15 @@ independent of this arg.)
 `score_scaling`<a id="score_scaling"></a>
 </td>
 <td>
-If true, the attention scores are divided by the square root
-of the dimension of keys (i.e., per_head_channels if transform_keys=True,
-else whatever the dimension of combined sender inputs is).
+One of either `"none"`, `"rsqrt_dim"`, or
+`"trainable_sigmoid"`. 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"`.
 </td>
 </tr><tr>
 <td>
@@ -270,13 +275,14 @@ the value transformation, then pools with attention coefficients.
 Setting this option pools inputs with attention coefficients, then applies
 the transformation. This is mathematically equivalent but can be faster
 or slower to compute, depending on the platform and the dataset.
-IMPORANT: Toggling this option breaks checkpoint compatibility.
+IMPORTANT: Toggling this option breaks checkpoint compatibility.
+IMPORTANT: Setting this option requires TensorFlow 2.10 or greater,
+because it uses `tf.keras.layers.EinsumDense`.
 </td>
 </tr>
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -335,7 +341,6 @@ Forwarded to the base class tf.keras.layers.Layer.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
@@ -368,7 +373,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 <h3 id="convolve"><code>convolve</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L332-L490">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L354-L530">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -394,7 +399,6 @@ 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`.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -482,7 +486,6 @@ does not require forwarding this arg, Keras does that automatically.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
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 8136ee3b..5c8df3d8 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L521-L574">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L561-L614">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -41,7 +41,6 @@ an edge set to do the analogous pooling of edge states to context.
 NOTE: This layer cannot pool node states. For that, use MultiHeadAttentionConv.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
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 a370f6ce..f5473cea 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L578-L640">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L618-L680">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -20,7 +20,7 @@ Returns a GraphUpdate layer with a transformer-style multihead attention.
     *,
     num_heads: int,
     per_head_channels: int,
-    receiver_tag: tfgnn.IncidentNodeOrContextTag,
+    receiver_tag: tfgnn.IncidentNodeTag,
     feature_name: str = tfgnn.HIDDEN_STATE,
     name: str = &#x27;multi_head_attention&#x27;,
     **kwargs
@@ -43,7 +43,6 @@ details).
 > itself requires having an explicit loop in the edge set.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
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 0f330280..05acc027 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L643-L736">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/layers.py#L683-L776">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -21,7 +21,7 @@ Returns a GraphUpdate layer for message passing with MultiHeadAttention pooling.
     units: int,
     message_dim: int,
     num_heads: int,
-    receiver_tag: tfgnn.IncidentNodeOrContextTag,
+    receiver_tag: tfgnn.IncidentNodeTag,
     node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None,
     edge_feature: Optional[tfgnn.FieldName] = None,
     l2_regularization: float = 0.0,
@@ -45,7 +45,6 @@ and all pooled messages, analogous to TF-GNN's
 `vanilla_mpnn.VanillaMPNNGraphUpdate` and `gat_v2.GATv2MPNNGraphUpdate`.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -159,7 +158,6 @@ Can be set to a `kerner_initializer` as understood by
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
diff --git a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/all_symbols.md b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/all_symbols.md
index 68cdf923..de2d267e 100644
--- a/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/all_symbols.md
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/all_symbols.md
@@ -9,3 +9,5 @@
 *   <a href="../multi_head_attention/MultiHeadAttentionEdgePool.md"><code>multi_head_attention.MultiHeadAttentionEdgePool</code></a>
 *   <a href="../multi_head_attention/MultiHeadAttentionHomGraphUpdate.md"><code>multi_head_attention.MultiHeadAttentionHomGraphUpdate</code></a>
 *   <a href="../multi_head_attention/MultiHeadAttentionMPNNGraphUpdate.md"><code>multi_head_attention.MultiHeadAttentionMPNNGraphUpdate</code></a>
+*   <a href="../multi_head_attention/graph_update_from_config_dict.md"><code>multi_head_attention.graph_update_from_config_dict</code></a>
+*   <a href="../multi_head_attention/graph_update_get_config_dict.md"><code>multi_head_attention.graph_update_get_config_dict</code></a>
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
new file mode 100644
index 00000000..ffa05192
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_from_config_dict.md
@@ -0,0 +1,74 @@
+# multi_head_attention.graph_update_from_config_dict
+
+[TOC]
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/config_dict.py#L40-L60">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Returns a MultiHeadAttentionMPNNGraphUpdate initialized from `cfg`.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>multi_head_attention.graph_update_from_config_dict(
+    cfg: config_dict.ConfigDict
+) -> tf.keras.layers.Layer
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`cfg`<a id="cfg"></a>
+</td>
+<td>
+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.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A new `MultiHeadAttentionMPNNGraphUpdate` object.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
+
+<tr>
+<td>
+`TypeError`<a id="TypeError"></a>
+</td>
+<td>
+if `cfg` fails to supply a required argument for
+`MultiHeadAttentionMPNNGraphUpdate.__init__`.
+</td>
+</tr>
+</table>
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
new file mode 100644
index 00000000..017fba0a
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/models/multi_head_attention/graph_update_get_config_dict.md
@@ -0,0 +1,22 @@
+# multi_head_attention.graph_update_get_config_dict
+
+[TOC]
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/models/multi_head_attention/config_dict.py#L22-L37">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Returns ConfigDict for graph_update_from_config_dict() with defaults.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>multi_head_attention.graph_update_get_config_dict() -> config_dict.ConfigDict
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
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 68df23df..1705f770 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
@@ -20,7 +20,7 @@ Returns a GraphUpdate layer for a Vanilla MPNN.
     *,
     units: int,
     message_dim: int,
-    receiver_tag: tfgnn.IncidentNodeOrContextTag,
+    receiver_tag: tfgnn.IncidentNodeTag,
     node_set_names: Optional[Collection[tfgnn.NodeSetName]] = None,
     edge_feature: Optional[tfgnn.FieldName] = None,
     reduce_type: str = &#x27;sum&#x27;,
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 2e62e5b1..b9f5be5f 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
@@ -22,8 +22,8 @@ Returns a VanillaMPNNGraphUpdate initialized from `cfg`.
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
-<!-- Tabular view -->
 
+<!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -44,7 +44,6 @@ passing it here.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
@@ -57,7 +56,6 @@ A new `VanillaMPNNGraphUpdate` object.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn.md b/tensorflow_gnn/docs/api_docs/python/tfgnn.md
index 4878c14f..59ef5ad7 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn.md
@@ -30,6 +30,9 @@ similar to `tf.TensorSpec`. For example, a `FieldSpec` describes an instance of
 
 ## Modules
 
+[`experimental`](./tfgnn/experimental.md) module: Experimental (unstable) parts
+of the public interface of TensorFlow GNN.
+
 [`keras`](./tfgnn/keras.md) module: The tfgnn.keras package.
 
 ## Classes
@@ -88,6 +91,9 @@ or from context to nodes or 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.
+
 [`check_homogeneous_graph_tensor(...)`](./tfgnn/check_homogeneous_graph_tensor.md):
 Raises ValueError unless there is exactly one node set and edge set.
 
@@ -134,6 +140,9 @@ dataset from generator of any nest of scalar graph pieces.
 [`mask_edges(...)`](./tfgnn/mask_edges.md): Creates a GraphTensor after applying
 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.
 
 [`parse_example(...)`](./tfgnn/parse_example.md): Parses a batch of serialized Example protos into a single `GraphTensor`.
@@ -271,7 +280,7 @@ TARGET_NAME<a id="TARGET_NAME"></a>
 **version**<a id="__version__"></a>
 </td>
 <td>
-`'0.4.0.dev1'`
+`'0.5.0.dev1'`
 </td>
 </tr>
 </table>
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md
index df244afe..c38a1771 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/Context.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L214-L328">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L215-L329">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -116,14 +116,14 @@ The total number of items.
 
 <h3 id="from_fields"><code>from_fields</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L225-L308">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L226-L309">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_fields(
     *,
-    features: Optional[<a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>] = None,
+    features: Optional[Fields] = None,
     sizes: Optional[Field] = None,
     shape: Optional[ShapeLike] = None,
     indices_dtype: Optional[tf.dtypes.DType] = None
@@ -206,7 +206,7 @@ A `Context` composite tensor.
 
 <h3 id="get_features_dict"><code>get_features_dict</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L157-L159">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L158-L160">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -218,12 +218,12 @@ Returns features copy as a dictionary.
 
 <h3 id="replace_features"><code>replace_features</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L310-L317">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L311-L318">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>replace_features(
-    features: <a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>
+    features: Fields
 ) -> 'Context'
 </code></pre>
 
@@ -246,13 +246,13 @@ Enforce the common prefix shape on all the contained features.
 
 <h3 id="__getitem__"><code>__getitem__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L54-L56">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L55-L57">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__getitem__(
     feature_name: FieldName
-) -> <a href="../tfgnn/Field.md"><code>tfgnn.Field</code></a>
+) -> Field
 </code></pre>
 
 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 34d9bd7e..583b0cfa 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/ContextSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/ContextSpec.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L331-L398">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L332-L399">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -116,14 +116,14 @@ Do NOT override for custom non-TF types.
 
 <h3 id="from_field_specs"><code>from_field_specs</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L335-L371">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L336-L372">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_field_specs(
     *,
-    features_spec: Optional[<a href="../tfgnn/FieldsSpec.md"><code>tfgnn.FieldsSpec</code></a>] = None,
+    features_spec: Optional[FieldsSpec] = None,
     sizes_spec: Optional[FieldSpec] = None,
     shape: ShapeLike = tf.TensorShape([]),
     indices_dtype: tf.dtypes.DType = const.default_indices_dtype
@@ -298,7 +298,7 @@ and `other`.
 
 <h3 id="relax"><code>relax</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L377-L398">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L378-L399">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -367,18 +367,15 @@ Return self==value.
 
 <h3 id="__getitem__"><code>__getitem__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L165-L166">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L166-L167">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__getitem__(
     feature_name: FieldName
-) -> <a href="../tfgnn/FieldSpec.md"><code>tfgnn.FieldSpec</code></a>
+) -> FieldSpec
 </code></pre>
 
-
-
-
 <h3 id="__ne__"><code>__ne__</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md
index 10ebf1b2..b5dab508 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSet.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L533-L616">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L534-L617">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -127,16 +127,13 @@ The total number of items.
 
 <h3 id="from_fields"><code>from_fields</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L556-L601">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L557-L602">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_fields(
-    *,
-    features: Optional[<a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>] = None,
-    sizes: <a href="../tfgnn/Field.md"><code>tfgnn.Field</code></a>,
-    adjacency: Adjacency
+    *, features: Optional[Fields] = None, sizes: Field, adjacency: Adjacency
 ) -> 'EdgeSet'
 </code></pre>
 
@@ -218,7 +215,7 @@ An `EdgeSet` composite tensor.
 
 <h3 id="get_features_dict"><code>get_features_dict</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L157-L159">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L158-L160">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -230,12 +227,12 @@ Returns features copy as a dictionary.
 
 <h3 id="replace_features"><code>replace_features</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L404-L410">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L405-L411">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>replace_features(
-    features: <a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>
+    features: Mapping[FieldName, Field]
 ) -> '_NodeOrEdgeSet'
 </code></pre>
 
@@ -258,13 +255,13 @@ Enforce the common prefix shape on all the contained features.
 
 <h3 id="__getitem__"><code>__getitem__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L54-L56">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L55-L57">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__getitem__(
     feature_name: FieldName
-) -> <a href="../tfgnn/Field.md"><code>tfgnn.Field</code></a>
+) -> Field
 </code></pre>
 
 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 e9076ed1..207e1a31 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSetSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/EdgeSetSpec.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L619-L678">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L620-L679">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -118,15 +118,15 @@ Do NOT override for custom non-TF types.
 
 <h3 id="from_field_specs"><code>from_field_specs</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L623-L637">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L624-L638">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_field_specs(
     *,
-    features_spec: Optional[<a href="../tfgnn/FieldsSpec.md"><code>tfgnn.FieldsSpec</code></a>] = None,
-    sizes_spec: <a href="../tfgnn/FieldSpec.md"><code>tfgnn.FieldSpec</code></a>,
+    features_spec: Optional[FieldsSpec] = None,
+    sizes_spec: FieldSpec,
     adjacency_spec: AdjacencySpec
 ) -> 'EdgeSetSpec'
 </code></pre>
@@ -299,7 +299,7 @@ and `other`.
 
 <h3 id="relax"><code>relax</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L653-L678">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L654-L679">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -376,18 +376,15 @@ Return self==value.
 
 <h3 id="__getitem__"><code>__getitem__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L165-L166">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L166-L167">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__getitem__(
     feature_name: FieldName
-) -> <a href="../tfgnn/FieldSpec.md"><code>tfgnn.FieldSpec</code></a>
+) -> FieldSpec
 </code></pre>
 
-
-
-
 <h3 id="__ne__"><code>__ne__</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md
index a70a7455..cf836449 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensor.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L681-L1188">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L682-L1189">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -236,13 +236,13 @@ The total number of graph components.
 
 <h3 id="from_pieces"><code>from_pieces</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L821-L853">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L822-L854">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_pieces(
-    context: Optional[<a href="../tfgnn/Context.md"><code>tfgnn.Context</code></a>] = None,
+    context: Optional[Context] = None,
     node_sets: Optional[Mapping[NodeSetName, NodeSet]] = None,
     edge_sets: Optional[Mapping[EdgeSetName, EdgeSet]] = None
 ) -> 'GraphTensor'
@@ -253,7 +253,7 @@ Constructs a new `GraphTensor` from context, node sets and edge sets.
 
 <h3 id="merge_batch_to_components"><code>merge_batch_to_components</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L855-L940">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L856-L941">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -334,7 +334,7 @@ A scalar (rank 0) graph tensor.
 
 <h3 id="remove_features"><code>remove_features</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1083-L1174">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1084-L1175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -447,12 +447,12 @@ input graph tensor.
 
 <h3 id="replace_features"><code>replace_features</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L982-L1081">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L983-L1082">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>replace_features(
-    context: Optional[<a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>] = None,
+    context: Optional[Fields] = None,
     node_sets: Optional[Mapping[NodeSetName, Fields]] = None,
     edge_sets: Optional[Mapping[EdgeSetName, Fields]] = None
 ) -> 'GraphTensor'
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md
index c6c7f938..141ade00 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/GraphTensorSpec.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1191-L1295">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1192-L1296">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -116,13 +116,13 @@ Do NOT override for custom non-TF types.
 
 <h3 id="from_piece_specs"><code>from_piece_specs</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1195-L1225">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1196-L1226">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_piece_specs(
-    context_spec: Optional[<a href="../tfgnn/ContextSpec.md"><code>tfgnn.ContextSpec</code></a>] = None,
+    context_spec: Optional[ContextSpec] = None,
     node_sets_spec: Optional[Mapping[NodeSetName, NodeSetSpec]] = None,
     edge_sets_spec: Optional[Mapping[EdgeSetName, EdgeSetSpec]] = None
 ) -> 'GraphTensorSpec'
@@ -296,7 +296,7 @@ and `other`.
 
 <h3 id="relax"><code>relax</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1258-L1295">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1259-L1296">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -322,7 +322,7 @@ Calling with all default parameters keeps the spec unchanged.
 `num_components`
 </td>
 <td>
-if True, allows the variable number of graph components.
+if True, allows a variable number of graph components.
 </td>
 </tr><tr>
 <td>
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md
index 7b6b3f60..55d71d06 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSet.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L418-L483">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L419-L484">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -124,15 +124,13 @@ The total number of items.
 
 <h3 id="from_fields"><code>from_fields</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L437-L474">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L438-L475">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_fields(
-    *,
-    features: Optional[<a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>] = None,
-    sizes: <a href="../tfgnn/Field.md"><code>tfgnn.Field</code></a>
+    *, features: Optional[Fields] = None, sizes: Field
 ) -> 'NodeSet'
 </code></pre>
 
@@ -202,7 +200,7 @@ A `NodeSet` composite tensor.
 
 <h3 id="get_features_dict"><code>get_features_dict</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L157-L159">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L158-L160">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -214,12 +212,12 @@ Returns features copy as a dictionary.
 
 <h3 id="replace_features"><code>replace_features</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L404-L410">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L405-L411">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>replace_features(
-    features: <a href="../tfgnn/Fields.md"><code>tfgnn.Fields</code></a>
+    features: Mapping[FieldName, Field]
 ) -> '_NodeOrEdgeSet'
 </code></pre>
 
@@ -242,13 +240,13 @@ Enforce the common prefix shape on all the contained features.
 
 <h3 id="__getitem__"><code>__getitem__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L54-L56">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L55-L57">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__getitem__(
     feature_name: FieldName
-) -> <a href="../tfgnn/Field.md"><code>tfgnn.Field</code></a>
+) -> Field
 </code></pre>
 
 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 151e3d89..d250edfe 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSetSpec.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/NodeSetSpec.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L486-L530">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L487-L531">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -116,15 +116,13 @@ Do NOT override for custom non-TF types.
 
 <h3 id="from_field_specs"><code>from_field_specs</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L490-L500">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L491-L501">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_field_specs(
-    *,
-    features_spec: Optional[<a href="../tfgnn/FieldsSpec.md"><code>tfgnn.FieldsSpec</code></a>] = None,
-    sizes_spec: <a href="../tfgnn/FieldSpec.md"><code>tfgnn.FieldSpec</code></a>
+    *, features_spec: Optional[FieldsSpec] = None, sizes_spec: FieldSpec
 ) -> 'NodeSetSpec'
 </code></pre>
 
@@ -296,7 +294,7 @@ and `other`.
 
 <h3 id="relax"><code>relax</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L506-L530">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L507-L531">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -373,18 +371,15 @@ Return self==value.
 
 <h3 id="__getitem__"><code>__getitem__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L165-L166">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L166-L167">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__getitem__(
     feature_name: FieldName
-) -> <a href="../tfgnn/FieldSpec.md"><code>tfgnn.FieldSpec</code></a>
+) -> FieldSpec
 </code></pre>
 
-
-
-
 <h3 id="__ne__"><code>__ne__</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
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 145a6aff..e2af57ae 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/all_symbols.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/all_symbols.md
@@ -36,6 +36,7 @@
 *   <a href="../tfgnn/broadcast_context_to_edges.md"><code>tfgnn.broadcast_context_to_edges</code></a>
 *   <a href="../tfgnn/broadcast_context_to_nodes.md"><code>tfgnn.broadcast_context_to_nodes</code></a>
 *   <a href="../tfgnn/broadcast_node_to_edges.md"><code>tfgnn.broadcast_node_to_edges</code></a>
+*   <a href="../tfgnn/check_compatible_with_schema_pb.md"><code>tfgnn.check_compatible_with_schema_pb</code></a>
 *   <a href="../tfgnn/check_homogeneous_graph_tensor.md"><code>tfgnn.check_homogeneous_graph_tensor</code></a>
 *   <a href="../tfgnn/check_required_features.md"><code>tfgnn.check_required_features</code></a>
 *   <a href="../tfgnn/check_scalar_graph_tensor.md"><code>tfgnn.check_scalar_graph_tensor</code></a>
@@ -44,6 +45,7 @@
 *   <a href="../tfgnn/create_schema_pb_from_graph_spec.md"><code>tfgnn.create_schema_pb_from_graph_spec</code></a>
 *   <a href="../tfgnn/dataset_filter_with_summary.md"><code>tfgnn.dataset_filter_with_summary</code></a>
 *   <a href="../tfgnn/dataset_from_generator.md"><code>tfgnn.dataset_from_generator</code></a>
+*   <a href="../tfgnn/experimental.md"><code>tfgnn.experimental</code></a>
 *   <a href="../tfgnn/find_tight_size_constraints.md"><code>tfgnn.find_tight_size_constraints</code></a>
 *   <a href="../tfgnn/gather_first_node.md"><code>tfgnn.gather_first_node</code></a>
 *   <a href="../tfgnn/get_io_spec.md"><code>tfgnn.get_io_spec</code></a>
@@ -63,6 +65,7 @@
 *   <a href="../tfgnn/keras/layers/ContextUpdate.md"><code>tfgnn.keras.layers.ContextUpdate</code></a>
 *   <a href="../tfgnn/keras/layers/EdgeSetUpdate.md"><code>tfgnn.keras.layers.EdgeSetUpdate</code></a>
 *   <a href="../tfgnn/keras/layers/GraphUpdate.md"><code>tfgnn.keras.layers.GraphUpdate</code></a>
+*   <a href="../tfgnn/keras/layers/ItemDropout.md"><code>tfgnn.keras.layers.ItemDropout</code></a>
 *   <a href="../tfgnn/keras/layers/MakeEmptyFeature.md"><code>tfgnn.keras.layers.MakeEmptyFeature</code></a>
 *   <a href="../tfgnn/keras/layers/MapFeatures.md"><code>tfgnn.keras.layers.MapFeatures</code></a>
 *   <a href="../tfgnn/keras/layers/NextStateFromConcat.md"><code>tfgnn.keras.layers.NextStateFromConcat</code></a>
@@ -78,6 +81,7 @@
 *   <a href="../tfgnn/keras/layers/SingleInputNextState.md"><code>tfgnn.keras.layers.SingleInputNextState</code></a>
 *   <a href="../tfgnn/learn_fit_or_skip_size_constraints.md"><code>tfgnn.learn_fit_or_skip_size_constraints</code></a>
 *   <a href="../tfgnn/mask_edges.md"><code>tfgnn.mask_edges</code></a>
+*   <a href="../tfgnn/node_degree.md"><code>tfgnn.node_degree</code></a>
 *   <a href="../tfgnn/pad_to_total_sizes.md"><code>tfgnn.pad_to_total_sizes</code></a>
 *   <a href="../tfgnn/parse_example.md"><code>tfgnn.parse_example</code></a>
 *   <a href="../tfgnn/parse_schema.md"><code>tfgnn.parse_schema</code></a>
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
new file mode 100644
index 00000000..6145c33e
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/check_compatible_with_schema_pb.md
@@ -0,0 +1,74 @@
+# tfgnn.check_compatible_with_schema_pb
+
+[TOC]
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/schema_utils.py#L188-L227">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Checks that the given spec or value is compatible with the graph schema.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfgnn.check_compatible_with_schema_pb(
+    graph: Union[<a href="../tfgnn/GraphTensor.md"><code>tfgnn.GraphTensor</code></a>, <a href="../tfgnn/GraphTensorSpec.md"><code>tfgnn.GraphTensorSpec</code></a>],
+    schema: <a href="../tfgnn/GraphSchema.md"><code>tfgnn.GraphSchema</code></a>
+) -> None
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+The `graph` is compatible with the `schema` if
+
+*   it is scalar (rank=0) graph tensor;
+*   has single graph component;
+*   has matching sets of nodes and edges;
+*   has matching sets of features on all node sets, edge sets, and the context,
+    and their types and shapes are compatible;
+*   all adjacencies are of type
+    <a href="../tfgnn/Adjacency.md"><code>tfgnn.Adjacency</code></a>.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`graph`<a id="graph"></a>
+</td>
+<td>
+The graph tensor or graph tensor spec.
+</td>
+</tr><tr>
+<td>
+`schema`<a id="schema"></a>
+</td>
+<td>
+The graph schema.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
+
+<tr>
+<td>
+`ValueError`<a id="ValueError"></a>
+</td>
+<td>
+if `spec_or_value` is not represented by the graph schema.
+</td>
+</tr>
+</table>
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 ca8c74a0..97c87e3a 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1539-L1546">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1540-L1547">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -17,8 +17,7 @@ Raises ValueError unless there is exactly one node set and edge set.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfgnn.check_homogeneous_graph_tensor(
-    graph: Union[<a href="../tfgnn/GraphTensor.md"><code>tfgnn.GraphTensor</code></a>, <a href="../tfgnn/GraphTensorSpec.md"><code>tfgnn.GraphTensorSpec</code></a>],
-    name=&#x27;This operation&#x27;
+    graph: Union[GraphTensor, GraphTensorSpec], name=&#x27;This operation&#x27;
 ) -> None
 </code></pre>
 
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 5a3fc187..40bd678e 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1412-L1414">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1413-L1415">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -15,11 +15,8 @@
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfgnn.check_scalar_graph_tensor(
-    graph: Union[<a href="../tfgnn/GraphTensor.md"><code>tfgnn.GraphTensor</code></a>, <a href="../tfgnn/GraphTensorSpec.md"><code>tfgnn.GraphTensorSpec</code></a>],
-    name=&#x27;This operation&#x27;
+    graph: Union[GraphTensor, GraphTensorSpec], name=&#x27;This operation&#x27;
 ) -> None
 </code></pre>
 
-
-
 <!-- Placeholder for "Used in" -->
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 74df00bf..66b57f8d 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
@@ -33,7 +33,8 @@ 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
-schema (as `tfgnn.check_compatible_with_schema_pb()`.)
+schema (as
+<a href="../tfgnn/check_compatible_with_schema_pb.md"><code>tfgnn.check_compatible_with_schema_pb()</code></a>.)
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
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 637e8bef..ac837446 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
@@ -30,10 +30,10 @@ 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
-(as `tfgnn.check_compatible_with_schema_pb()`.)
+(as
+<a href="../tfgnn/check_compatible_with_schema_pb.md"><code>tfgnn.check_compatible_with_schema_pb()</code></a>.)
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -49,7 +49,6 @@ The scalar graph tensor or its spec with single graph component.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
@@ -62,7 +61,6 @@ An instance of the graph schema proto message.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
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 842e571e..fa38d480 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
@@ -42,7 +42,6 @@ print([dataset1])  # prints: pieceA, pieceD.
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -61,7 +60,6 @@ protocol. Could consist of any nest of tensors and scalar graph pieces
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
@@ -74,7 +72,6 @@ A `tf.data.Dataset`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md
new file mode 100644
index 00000000..97dbe4ac
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/experimental.md
@@ -0,0 +1,31 @@
+# Module: tfgnn.experimental
+
+[TOC]
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/experimental/__init__.py">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Experimental (unstable) parts of the public interface of TensorFlow GNN.
+
+A symbol `foo` exposed here is available to library users as
+
+```
+import tensorflow_gnn as tfgnn
+
+tfgnn.experimental.foo()
+```
+
+This is the preferred way to expose individual functions on track to inclusion
+into the stable public interface of TensorFlow GNN.
+
+Beyond these symbols, there are also experimental sub-libraries that need to be
+imported separately (`from tensorflow_gnn.experimental import foo`). This is for
+special cases only.
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md
index bd3e8518..12b66418 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/homogeneous.md
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1457-L1536">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor.py#L1458-L1537">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -27,13 +27,12 @@ Constructs a homogeneous `GraphTensor` with node features and one edge_set.
     edge_set_name: Optional[FieldName] = const.EDGES,
     node_set_sizes: Optional[Field] = None,
     edge_set_sizes: Optional[Field] = None
-) -> <a href="../tfgnn/GraphTensor.md"><code>tfgnn.GraphTensor</code></a>
+) -> GraphTensor
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
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 d66b391e..d534e1e6 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/tensor_utils.py#L478-L480">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/tensor_utils.py#L486-L488">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
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 90716515..47224171 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/ConvGNNBuilder.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/ConvGNNBuilder.md
@@ -20,12 +20,10 @@ Factory of layers that do convolutions on a graph.
     convolutions_factory: Callable[..., graph_update_lib.EdgesToNodePoolingLayer],
     nodes_next_state_factory: Callable[[const.NodeSetName], next_state_lib.NextStateForNodeSet],
     *,
-    receiver_tag: Optional[const.IncidentNodeOrContextTag] = None
+    receiver_tag: Optional[const.IncidentNodeTag] = None
 )
 </code></pre>
 
-
-
 <!-- Placeholder for "Used in" -->
 
 ConvGNNBuilder object constructs `GraphUpdate` layers, that apply arbitrary
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 60206410..a6547d36 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers.md
@@ -34,6 +34,9 @@ self-loops to scalar graphs.
 
 [`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 MapFeatures`](../../tfgnn/keras/layers/MapFeatures.md): Transforms features on a GraphTensor by user-defined callbacks.
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 4203105e..f6ae3e81 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
@@ -91,7 +91,6 @@ To pass the default state tensor of the context, set this to
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Call returns</h2></th></tr>
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
new file mode 100644
index 00000000..e3132f55
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/keras/layers/ItemDropout.md
@@ -0,0 +1,87 @@
+# tfgnn.keras.layers.ItemDropout
+
+[TOC]
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/item_dropout.py#L22-L74">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Dropout of feature values for entire edges, nodes or components.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfgnn.keras.layers.ItemDropout(
+    rate: float, seed: Optional[int] = None, **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+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.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Init args</h2></th></tr>
+
+<tr>
+<td>
+`rate`<a id="rate"></a>
+</td>
+<td>
+The dropout rate, forwarded to `tf.keras.layers.Dropout`.
+</td>
+</tr><tr>
+<td>
+`seed`<a id="seed"></a>
+</td>
+<td>
+The random seed, forwarded `tf.keras.layers.Dropout`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Call args</h2></th></tr>
+
+<tr>
+<td>
+`x`<a id="x"></a>
+</td>
+<td>
+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.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Call returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+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.
+</td>
+</tr>
+
+</table>
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 51ddf270..242b606a 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
@@ -63,7 +63,6 @@ a Context, NodeSet or EdgeSet from a GraphTensor.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Call returns</h2></th></tr>
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 8342a854..6de78c2f 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
@@ -47,10 +47,8 @@ unset to select tfgnn.HIDDEN_STATE.
 
 #### For example:
 
-
-
 ```python
-readout = tfgnn.keras.layers.Readout(feature="value")
+readout = tfgnn.keras.layers.Readout(feature_name="value")
 value = readout(graph_tensor, edge_set_name="edges")
 assert value == graph_tensor.edge_sets["edge"]["value"]
 ```
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 c846ebe8..1145f376 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/next_state.py#L145-L228">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/next_state.py#L145-L238">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -30,10 +30,15 @@ Updates a state with a residual block.
 <!-- Placeholder for "Used in" -->
 
 This layer concatenates all inputs, sends them through a user-supplied
-transformation, forms a skip connection by adding back the state of the
-updated graph piece, and finally applies an activation function.
-In other words, the user-supplied transformation is a residual block
-that modifies the state.
+transformation, forms a skip connection by adding back the state of the updated
+graph piece, and finally applies an activation function. In other words, the
+user-supplied transformation is a residual block that modifies the state. The
+output shape of the residual block must match the shape of the state that gets
+updated so that they can be added.
+
+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.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
@@ -46,18 +51,18 @@ that modifies the state.
 </td>
 <td>
 Required. A Keras Layer to transform the concatenation
-of all inputs into a delta that gets added to the state. Notice that
-the activation function is applied after the residual_block and the
-addition, so typically the residual_block does *not* use an activation
-function in its last layer.
+of all inputs into a delta that gets added to the state.
 </td>
 </tr><tr>
 <td>
 `activation`<a id="activation"></a>
 </td>
 <td>
-An activation function (none by default),
-as understood by tf.keras.layers.Activation.
+An activation function (none by default), as understood by
+`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.
 </td>
 </tr><tr>
 <td>
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 3aff814e..37370c07 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/convolutions.py#L26-L142">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/convolutions.py#L26-L144">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -245,7 +245,7 @@ If `False`, all calls to convolve() will get `sender_node_input=None`.
 
 <h3 id="convolve"><code>convolve</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/convolutions.py#L120-L142">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/convolutions.py#L120-L144">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -257,6 +257,7 @@ source</a>
     broadcast_from_sender_node: Callable[[tf.Tensor], tf.Tensor],
     broadcast_from_receiver: Callable[[tf.Tensor], tf.Tensor],
     pool_to_receiver: Callable[..., tf.Tensor],
+    extra_receiver_ops: Any = None,
     training: bool
 ) -> tf.Tensor
 </code></pre>
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 8388d80d..efb8efaf 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
@@ -6,7 +6,7 @@
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/next_state.py#L231-L258">
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/keras/layers/next_state.py#L241-L268">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
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 e60b4462..a385fa63 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/mask_edges.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/mask_edges.md
@@ -33,7 +33,6 @@ Edge masking doesn't change the node sets or the context node information.
 Not compatible with XLA.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -72,7 +71,6 @@ new edge-set, with name masked_info_edge_set_name.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
diff --git a/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md b/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md
new file mode 100644
index 00000000..31dbc255
--- /dev/null
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/node_degree.md
@@ -0,0 +1,74 @@
+# tfgnn.node_degree
+
+[TOC]
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/gnn/tree/master/tensorflow_gnn/graph/graph_tensor_ops.py#L1173-L1200">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Returns the degree of each node w.r.t. one side of an edge set.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfgnn.node_degree(
+    graph_tensor: <a href="../tfgnn/GraphTensor.md"><code>tfgnn.GraphTensor</code></a>,
+    edge_set_name: EdgeSetName,
+    node_tag: IncidentNodeTag
+) -> <a href="../tfgnn/Field.md"><code>tfgnn.Field</code></a>
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`graph_tensor`<a id="graph_tensor"></a>
+</td>
+<td>
+A scalar GraphTensor.
+</td>
+</tr><tr>
+<td>
+`edge_set_name`<a id="edge_set_name"></a>
+</td>
+<td>
+The name of the edge set for which degrees are calculated.
+</td>
+</tr><tr>
+<td>
+`node_tag`<a id="node_tag"></a>
+</td>
+<td>
+The side of each edge for which the degrees are calculated,
+specified by its tag in the edge set (e.g., <a href="../tfgnn.md#SOURCE"><code>tfgnn.SOURCE</code></a>,
+<a href="../tfgnn.md#TARGET"><code>tfgnn.TARGET</code></a>).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+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.
+</td>
+</tr>
+
+</table>
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 57d22232..9e09763d 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/reorder_nodes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/reorder_nodes.md
@@ -25,8 +25,8 @@ Reorders nodes within node sets according to indices.
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
-<!-- Tabular view -->
 
+<!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -59,7 +59,6 @@ If True, checks that `node_indices` are valid permutations.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
@@ -72,7 +71,6 @@ A scalar GraphTensor with randomly shuffled nodes within `node_sets`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
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 33e82e7a..ae032b9b 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
@@ -24,6 +24,7 @@ Shuffles context, node set and edge set features of a scalar GraphTensor.
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
+
 <!-- Tabular view -->
 
  <table class="responsive fixed orange">
@@ -48,7 +49,6 @@ A seed for random uniform shuffle.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
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 0edd0c1e..9b02a7bc 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_nodes.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/shuffle_nodes.md
@@ -32,7 +32,6 @@ the new order of shuffled nodes. The order of graph components (as created by
 within each component.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -63,7 +62,6 @@ A seed for random uniform shuffle.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
@@ -76,7 +74,6 @@ A scalar GraphTensor with randomly shuffled nodes within `node_sets`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
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 9dfffd05..e2279ed5 100644
--- a/tensorflow_gnn/docs/api_docs/python/tfgnn/write_example.md
+++ b/tensorflow_gnn/docs/api_docs/python/tfgnn/write_example.md
@@ -31,9 +31,10 @@ Python job. Create instances of scalar `GraphTensor` with a single graph
 component and call this to write them out. It is recommended to always accompany
 serialized graph tensor tensorflow examples by their graph schema file (see
 <a href="../tfgnn/create_schema_pb_from_graph_spec.md"><code>tfgnn.create_schema_pb_from_graph_spec()</code></a>).
-TF-GNN library provides `tfgnn.check_compatible_with_schema_pb()` to check that
-graph tensor instances (or their specs) are compatible with the graph schema.
-The graph tensors materialized in this way will be parseable by
+TF-GNN library provides
+<a href="../tfgnn/check_compatible_with_schema_pb.md"><code>tfgnn.check_compatible_with_schema_pb()</code></a>
+to check that graph tensor instances (or their specs) are compatible with the
+graph schema. The graph tensors materialized in this way will be parseable by
 <a href="../tfgnn/parse_example.md"><code>tfgnn.parse_example()</code></a>
 (using the spec deserialized from the schema) and have the same contents (up to
 the choice of indices_dtype).