Library updates. (#131)

Pushing updates to the library:

* Added support for nested features in utils_tf.concat.
* Added support for padding graphs to fix sizes for TPU compatibility.
* Updated documentation.
* Updated README.

README.md

@@ -1,4 +1,4 @@
-![Graph Nets DeepMind shortest path](images/graph-nets-deepmind-shortest-path0.gif)
+![Graph Nets DeepMind shortest path](https://github.com/deepmind/graph_nets/raw/master/images/graph-nets-deepmind-shortest-path0.gif)
 
 # Graph Nets library
 
@@ -17,7 +17,7 @@ part of the broader family of "graph neural networks" (Scarselli et al., 2009).
 To learn more about graph networks, see our arXiv paper: [Relational inductive
 biases, deep learning, and graph networks](https://arxiv.org/abs/1806.01261).
 
-![Graph network](images/graph-network.png)
+![Graph network](https://github.com/deepmind/graph_nets/raw/master/images/graph-network.png)
 
 ## Installation
 
@@ -99,7 +99,7 @@ label the nodes and edges on the shortest path between any two nodes. Over a
 sequence of message-passing steps (as depicted by each step's plot), the
 model refines its prediction of the shortest path.
 
-![Shortest path](images/shortest-path.png)
+![Shortest path](https://github.com/deepmind/graph_nets/raw/master/images/shortest-path.png)
 
 ----------------
 
@@ -110,7 +110,7 @@ sort the list. After a sequence of message-passing steps, the model makes an
 accurate prediction of which elements (columns in the figure) come next after
 each other (rows).
 
-![Sort](images/sort.png)
+![Sort](https://github.com/deepmind/graph_nets/raw/master/images/sort.png)
 
 ----------------
 
@@ -123,11 +123,11 @@ a future trajectory. Each subplot below shows the true and predicted mass-spring
 system states over 50 steps. This is similar to the model and experiments in
 Battaglia et al. (2016)'s "interaction networks".
 
-![Physics](images/physics.png)
+![Physics](https://github.com/deepmind/graph_nets/raw/master/images/physics.png)
 
 ----------------
 
-#### [Run "graph nets basics demo" in browser](https://colab.research.google.com/github/deepmind/graph_nets/blob/master/graph_nets/demos/graph_nets_basics.ipynb)
+#### [Run "graph nets basics demo" in browser](https://colab.research.google.com/github/deepmind/graph_nets/blob/master/graph_nets/demos/graph_nets_basics.ipynb)  [(Run TF2 version)](https://colab.research.google.com/github/deepmind/graph_nets/blob/master/graph_nets/demos_tf2/graph_nets_basics.ipynb)
 
 The "graph nets basics demo" is a tutorial containing step by step examples
 about how to create and manipulate graphs, how to feed them into
@@ -154,11 +154,13 @@ then open a demo through the Jupyter notebook interface.
 
 ## Other graph neural network libraries
 
-If you use PyTorch, check out these high-quality open-source libraries for
-graph neural networks:
+Check out these high-quality open-source libraries for graph neural networks:
+
+* [jraph](https://github.com/deepmind/jraph): DeepMind's GNNs/GraphNets library
+for [JAX](https://github.com/google/jax).
 
 * [pytorch_geometric](https://github.com/rusty1s/pytorch_geometric): See
 [MetaLayer](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.meta.MetaLayer)
 for an analog of our Graph Nets interface.
 
-* [Deep Graph Library (DGL)](https://github.com/dmlc/dgl)
+* [Deep Graph Library (DGL)](https://github.com/dmlc/dgl).

docs/contents.md

@@ -38,12 +38,17 @@
   * [`utils_tf.fully_connect_graph_static`](graph_nets.md#utils_tf.fully_connect_graph_static)
   * [`utils_tf.get_feed_dict`](graph_nets.md#utils_tf.get_feed_dict)
   * [`utils_tf.get_graph`](graph_nets.md#utils_tf.get_graph)
+  * [`utils_tf.get_graphs_tuple_size`](graph_nets.md#utils_tf.get_graphs_tuple_size)
+  * [`utils_tf.get_mask`](graph_nets.md#utils_tf.get_mask)
   * [`utils_tf.get_num_graphs`](graph_nets.md#utils_tf.get_num_graphs)
+  * [`utils_tf.GraphsTupleSize`](graph_nets.md#utils_tf.GraphsTupleSize)
   * [`utils_tf.identity`](graph_nets.md#utils_tf.identity)
   * [`utils_tf.make_runnable_in_session`](graph_nets.md#utils_tf.make_runnable_in_session)
   * [`utils_tf.nest_to_numpy`](graph_nets.md#utils_tf.nest_to_numpy)
+  * [`utils_tf.pad_graphs_tuple`](graph_nets.md#utils_tf.pad_graphs_tuple)
   * [`utils_tf.placeholders_from_data_dicts`](graph_nets.md#utils_tf.placeholders_from_data_dicts)
   * [`utils_tf.placeholders_from_networkxs`](graph_nets.md#utils_tf.placeholders_from_networkxs)
+  * [`utils_tf.remove_graphs_tuple_padding`](graph_nets.md#utils_tf.remove_graphs_tuple_padding)
   * [`utils_tf.repeat`](graph_nets.md#utils_tf.repeat)
   * [`utils_tf.set_zero_edge_features`](graph_nets.md#utils_tf.set_zero_edge_features)
   * [`utils_tf.set_zero_global_features`](graph_nets.md#utils_tf.set_zero_global_features)

graph_nets/_base.py

@@ -67,7 +67,7 @@ class AbstractModule(snt.Module):
 
     def __init__(self, *args, **kwargs):
       super(AbstractModule, self).__init__(*args, **kwargs)
-      self.__call__.__func__.__doc__ = self._build.__doc__
+      self.__call__.__func__.__doc__ = self._build.__doc__  # pytype: disable=attribute-error
 
     # In snt2 calls to `_enter_variable_scope` are ignored.
     @contextlib.contextmanager