2026a8939d9ccc3e26ac776db5b4788846fd166c,dnc/sparse_memory.py,SparseMemory,write_into_sparse_memory,#SparseMemory#Any#,124

Before Change

    pos = positions.data.cpu().numpy()
    for b in range(positions.size(0)):
      // update sparse memory
      for i, p in enumerate(list(pos[b])):
        hidden["memory"][b, p] = read_vectors[b, i]

      // update indexes
      hidden["indexes"][b].add(read_vectors[b], positions[b])
      hidden["last_used_mem"][b] = (int(pos[b][-1]) + 1) if (pos[b][-1] + 1) < self.mem_size else 0

    return hidden

After Change

    (b, m, w) = hidden["memory"].size()
    // update memory
    hidden["memory"].scatter_(1, positions.unsqueeze().expand(b, self.K+1, w), read_vectors)

    // non-differentiable operations
    pos = positions.data.cpu().numpy()

Instances