Calculate max along one dimension in tensorflow tensor

Question

I have tf tensor in the form of [number_of_image, width, height, channel] . The channel dim is optional and can be removed. I would like to calulate max value for each image. It should be as fast as possible and should work in graphic mode of tensorflow execution.

Max calculation is for max normalization of each image. I tried to use tf.reduce_max() with axis=0 option but it gives me tensor of with size of [width, height, channel] which is weird. I ended up with unstacking and stacking (code below) but I wonder if there is a better and fast soluton?

#grad is tensor with form [number_of_image, width, height, channel]
grad_unpack = tf.unstack(grad)
for t in grad_unpack:
    t /= tf.reduce_max(t)    
grad = tf.stack(grad_unpack)

TIA

Answer 1

tf.reduce_max(grad, axis=[1, 2, 3])

For any rank tensor is could be more tricky

tf.reduce_max(tf.reshape(grad, shape=[tf.shape(grad)[0], -1]), 1)