it depends how your AI is implemented, perhaps it will turn out that the first human level AI's are simply massive ANN's of some kind. Such an AI might have human equivalent working memory and have to do the equivalent of making notes outside of it's own mind just as we do.

Given how very very far we are from that level of AI we might very well see actual brain enhancements similar to this only for humans first which could leave us on a much for even footing with the AI's:

http://www.popsci.com/technology/article/2011-06/artificial-memory-chip-rats-can-remember-and-forget-touch-button

The device can mimic the brain's own neural signals, thereby serving as a surrogate for a piece of the brain associated with forming memories. If there is sufficient neural activity to trace, the device can restore memories after they have been lost. If it's used with a normal, functioning hippocampus, the device can even enhance memory.