memory.py 文件源码

def make_update_op(self, upd_idxs, upd_keys, upd_vals,
                       batch_size, use_recent_idx, intended_output):
        """Function that creates all the update ops."""
        base_update_op = super(LSHMemory, self).make_update_op(
            upd_idxs, upd_keys, upd_vals,
            batch_size, use_recent_idx, intended_output)

        hash_slot_idxs = self.get_hash_slots(upd_keys)

        update_ops = []
        with tf.control_dependencies([base_update_op]):
            for i, slot_idxs in enumerate(hash_slot_idxs):
                entry_idx = tf.random_uniform([batch_size],
                                              maxval=self.num_per_hash_slot,
                                              dtype=tf.int32)
                entry_mul = 1 - tf.one_hot(entry_idx, self.num_per_hash_slot,
                                           dtype=tf.int32)
                entry_add = (tf.expand_dims(upd_idxs, 1) *
                             tf.one_hot(entry_idx, self.num_per_hash_slot,
                                        dtype=tf.int32))

                mul_op = tf.scatter_mul(
                    self.hash_slots[i], slot_idxs, entry_mul)
                with tf.control_dependencies([mul_op]):
                    add_op = tf.scatter_add(
                        self.hash_slots[i], slot_idxs, entry_add)
                    update_ops.append(add_op)

        return tf.group(*update_ops)