Operating over vmap array

[dingjr7]

Dear communities,

I am trying to implement quantization in a vmap function. However, what i need to know is the min max of the batched array, not an single element. My current implementation is as follow,

xmin = lax.pmin(x, axis_name=axis_name)
xmax = lax.pmax(x, axis_name=axis_name)

Although it is hidden behind lax.stop_gradient(), it fails to compile because "NotImplementedError: Differentiation rule for 'pmin' not implemented".

Best regards!