(x, batch)
| 18 | |
| 19 | |
| 20 | def to_dense_batch(x, batch): |
| 21 | assert batch is not None |
| 22 | batch_size = tf.reduce_max(batch) + 1 |
| 23 | num_nodes = mp_ops.scatter_('add', tf.ones([tf.shape(batch)[0], 1]), |
| 24 | batch, batch_size) |
| 25 | num_nodes = tf.cast(tf.reshape(num_nodes, [-1]), dtype=tf.int32) |
| 26 | |
| 27 | cum_nodes = tf.concat([tf.zeros(1, dtype=tf.int32), |
| 28 | tf.cumsum(num_nodes, axis=0)], axis=0) |
| 29 | max_num_nodes = tf.reduce_max(num_nodes) |
| 30 | |
| 31 | idx = tf.range(tf.reduce_sum(num_nodes)) |
| 32 | |
| 33 | n = tf.gather(cum_nodes, batch) |
| 34 | idx = idx - n + batch * max_num_nodes |
| 35 | |
| 36 | idx = tf.reshape(idx, [-1, 1]) |
| 37 | |
| 38 | size = [batch_size * max_num_nodes, tf.shape(x)[-1]] |
| 39 | |
| 40 | out = tf.scatter_nd(idx, x, shape=size) |
| 41 | |
| 42 | out_size = [batch_size, max_num_nodes, tf.shape(x)[-1]] |
| 43 | out = tf.reshape(out, out_size) |
| 44 | |
| 45 | mask = tf.scatter_nd(idx, |
| 46 | tf.ones(tf.shape(batch)[0]), |
| 47 | shape=[batch_size * max_num_nodes]) |
| 48 | return out, out_size, mask |
no outgoing calls