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))