xla

Fix output tensor shape for argmin and argmax where keepdim=True and dim=None

#6536

Merged