| 728 | return tf.group(*update_ops) |
| 729 | |
| 730 | def _optimize_store(self, node, node_embeddings): |
| 731 | if not self.gradient_stores: |
| 732 | return tf.zeros([]), tf.no_op() |
| 733 | |
| 734 | losses = [] |
| 735 | clear_ops = [] |
| 736 | for gradient_store, node_embedding in zip(self.gradient_stores, |
| 737 | node_embeddings): |
| 738 | embedding_gradient = tf.nn.embedding_lookup(gradient_store, node) |
| 739 | with tf.control_dependencies([embedding_gradient]): |
| 740 | clear_ops.append( |
| 741 | utils_embedding.embedding_update( |
| 742 | gradient_store, node, |
| 743 | tf.zeros_like(embedding_gradient))) |
| 744 | losses.append(tf.reduce_sum(node_embedding * embedding_gradient)) |
| 745 | |
| 746 | store_loss = tf.add_n(losses) |
| 747 | with tf.control_dependencies(clear_ops): |
| 748 | return store_loss, self.store_optimizer.minimize(store_loss) |
| 749 | |
| 750 | |
| 751 | class LayerEncoder(SageEncoder): |