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.
Edited by Gordian Edenhofer
JAX deprecated ndarray.device
in favor of ndarray.devices
(with an s
) returning a list of devices instead of a single device.