lmap: Better determine like device using un-deprecated method

JAX deprecated ndarray.device in favor of ndarray.devices (with an s) returning a list of devices instead of a single device.