title: Tensorflow multi-dimension argmax
body: <p>Say I have a tensor of size BxWxHxD. I want to process the tensor such that I have a new BxWxHxD tensor where only the maximum element in each WxH slice is kept, and all other values are zero.</p>
  
  <p>In other words, I think the best way to achieve this is to somehow take a 2D argmax across the WxH slices, resulting in BxD index tensors for the rows and colums that can then be converted to a one-hot BxWxHxD tensor to be used as a mask. How do I make this work?</p>