(self, x, edge_index, size=None, **kwargs)
| 37 | self.rnn = tf.nn.rnn_cell.MultiRNNCell(lstm_cells) |
| 38 | |
| 39 | def __call__(self, x, edge_index, size=None, **kwargs): |
| 40 | h = x |
| 41 | for i in range(self.processing_steps): |
| 42 | m = [None if h[0] is None else self.fc[i](h[0]), |
| 43 | None if h[1] is None else self.fc[i](h[1])] |
| 44 | gather_x, = self.gather_feature([m], edge_index) |
| 45 | out = self.apply_edge(gather_x[1]) |
| 46 | out = mp_ops.scatter_(self.aggr, out, edge_index[0], size=size[0]) |
| 47 | out = self.apply_node(out) |
| 48 | out = tf.expand_dims(out, axis=1) |
| 49 | hidden_state = [h[0] for _ in range(self.lstm_layers)] |
| 50 | with tf.variable_scope('rnn', reuse=tf.AUTO_REUSE): |
| 51 | out, _ = tf.nn.dynamic_rnn(self.rnn, |
| 52 | out, |
| 53 | initial_state=tuple(hidden_state), |
| 54 | dtype=tf.float32) |
| 55 | out = tf.reshape(out, [-1, self.dim]) |
| 56 | h = [out, h[1]] |
| 57 | return out |
| 58 | |
| 59 | def apply_edge(self, x_j): |
| 60 | return x_j |
nothing calls this directly
no test coverage detected