keras/src/backend/torch/numpy.py/amax
def amax(x, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    if axis is None:
        return torch.amax(x)
    if axis == () or axis == []:
        # Torch handles the empty axis case differently from numpy.
        return x
    return torch.amax(x, dim=axis, keepdim=keepdims)