def _max_mask_non_finite(x, axis=-1, keepdims=False, mask=0):
  """"""
  m = np.max(x, axis=_astuple(axis), keepdims=keepdims)
  needs_masking = ~np.isfinite(m)
  if needs_masking.ndim > 0:
    m[needs_masking] = mask
  elif needs_masking:
    m = mask
  return m