def test_compute_stacked_offsets(self):
offset0 = utils_tf._compute_stacked_offsets(self.sizes, self.repeats)
offset1 = utils_tf._compute_stacked_offsets(np.array(self.sizes),
np.array(self.repeats))
offset2 = utils_tf._compute_stacked_offsets(
tf.constant(self.sizes, dtype=tf.int32),
tf.constant(self.repeats, dtype=tf.int32))
self.assertAllEqual(self.offset, offset0.numpy().tolist())
self.assertAllEqual(self.offset, offset1.numpy().tolist())
self.assertAllEqual(self.offset, offset2.numpy().tolist())
def test_compute_stacked_offsets(self):
offset0 = utils_tf._compute_stacked_offsets(self.sizes, self.repeats)
offset1 = utils_tf._compute_stacked_offsets(
np.array(self.sizes), np.array(self.repeats))
offset2 = utils_tf._compute_stacked_offsets(
tf.constant(self.sizes, dtype=tf.int32),
tf.constant(self.repeats, dtype=tf.int32))
with self.test_session() as sess:
o0, o1, o2 = sess.run([offset0, offset1, offset2])
self.assertAllEqual(self.offset, o0.tolist())
self.assertAllEqual(self.offset, o1.tolist())
self.assertAllEqual(self.offset, o2.tolist())
def autoregressive_connect_graph_dynamic(
graph,
exclude_self_edges=False,
name="autoregressive_connect_graph_dynamic"):
"""Adds edges to a graph by fully-connecting the nodes.
This method does not require the number of nodes per graph to be constant,
or to be known at graph building time.
Args:
graph: A `graphs.GraphsTuple` with `None` values for the edges, senders and
receivers.
exclude_self_edges (default=False): Excludes self-connected edges.
name: (string, optional) A name for the operation.
Returns:
A `graphs.GraphsTuple` containing `Tensor`s with fully-connected edges.
Raises:
ValueError: if any of the `EDGES`, `RECEIVERS` or `SENDERS` field is not
`None` in `graph`.
"""
utils_tf._validate_edge_fields_are_all_none(graph)
with tf.name_scope(name):
def body(i, senders, receivers, n_edge):
edges = _create_autogressive_edges_from_nodes_dynamic(
graph.n_node[i], exclude_self_edges)
return (i + 1, senders.write(i, edges['senders']),
receivers.write(i, edges['receivers']),
n_edge.write(i, edges['n_edge']))
num_graphs = utils_tf.get_num_graphs(graph)
loop_condition = lambda i, *_: tf.less(i, num_graphs)
initial_loop_vars = [0] + [
tf.TensorArray(dtype=tf.int32, size=num_graphs, infer_shape=False)
for _ in range(3) # senders, receivers, n_edge
]
_, senders_array, receivers_array, n_edge_array = tf.while_loop(
loop_condition, body, initial_loop_vars, back_prop=False)
n_edge = n_edge_array.concat()
offsets = utils_tf._compute_stacked_offsets(graph.n_node, n_edge)
senders = senders_array.concat() + offsets
receivers = receivers_array.concat() + offsets
senders.set_shape(offsets.shape)
receivers.set_shape(offsets.shape)
receivers.set_shape([None])
senders.set_shape([None])
num_graphs = graph.n_node.get_shape().as_list()[0]
n_edge.set_shape([num_graphs])
return graph.replace(senders=senders,
receivers=receivers,
n_edge=n_edge)
def connect_graph_dynamic(graph: GraphsTuple, is_edge_func, name="connect_graph_dynamic"):
"""
Connects a graph using a boolean edge mask to create edges.
Args:
graph: GraphsTuple
is_edge_func: callable(sender: int, receiver: int) -> bool, should broadcast
name:
Returns:
connected GraphsTuple
"""
utils_tf._validate_edge_fields_are_all_none(graph)
with tf.name_scope(name):
def body(i, senders, receivers, n_edge):
edges = _create_functional_connect_edges_dynamic(graph.n_node[i], is_edge_func)
# edges = create_edges_func(graph.n_node[i])
return (i + 1, senders.write(i, edges['senders']),
receivers.write(i, edges['receivers']),
n_edge.write(i, edges['n_edge']))
num_graphs = utils_tf.get_num_graphs(graph)
loop_condition = lambda i, *_: tf.less(i, num_graphs)
initial_loop_vars = [0] + [
tf.TensorArray(dtype=tf.int32, size=num_graphs, infer_shape=False)
for _ in range(3) # senders, receivers, n_edge
]
_, senders_array, receivers_array, n_edge_array = tf.while_loop(loop_condition, body, initial_loop_vars)
n_edge = n_edge_array.concat()
offsets = utils_tf._compute_stacked_offsets(graph.n_node, n_edge)
senders = senders_array.concat() + offsets
receivers = receivers_array.concat() + offsets
senders.set_shape(offsets.shape)
receivers.set_shape(offsets.shape)
receivers.set_shape([None])
senders.set_shape([None])
num_graphs = graph.n_node.get_shape().as_list()[0]
n_edge.set_shape([num_graphs])
return graph.replace(senders=tf.stop_gradient(senders),
receivers=tf.stop_gradient(receivers),
n_edge=tf.stop_gradient(n_edge))