How does argamx work when given a 3d tensor - tensorflow

Question

I was wondering how does argmax work when given a 3D tensor. I know what happens when it has a 2D tesnor but 3D is confusing me a lot.

Example:

import