(self, inputs, index, size=None)
| 31 | self.lstm = tf.nn.rnn_cell.MultiRNNCell(lstm_cells) |
| 32 | |
| 33 | def __call__(self, inputs, index, size=None): |
| 34 | size = tf.reduce_max(index) + 1 if size is None else size |
| 35 | cell_in = tf.zeros([size, self.dim*2], dtype=tf.float32) |
| 36 | hidden_state = self.lstm.zero_state(tf.shape(cell_in)[0], |
| 37 | dtype=tf.float32) |
| 38 | for i in range(self.processing_steps): |
| 39 | q = tf.expand_dims(cell_in, axis=1) |
| 40 | q, hidden_state = tf.nn.dynamic_rnn(self.lstm, |
| 41 | q, |
| 42 | initial_state=hidden_state, |
| 43 | dtype=tf.float32) |
| 44 | q = tf.reshape(q, [-1, self.dim]) |
| 45 | e = tf.reduce_sum((inputs * tf.gather(q, index)), |
| 46 | axis=-1, |
| 47 | keep_dims=True) |
| 48 | a = mp_ops.scatter_softmax(e, index, size=size) |
| 49 | r = mp_ops.scatter_(self.aggr, a * inputs, index, size=size) |
| 50 | cell_in = tf.reshape(tf.concat([q, r], axis=-1), [-1, self.dim*2]) |
| 51 | return cell_in |
nothing calls this directly
no outgoing calls
no test coverage detected