Memstats are not guaranteed to be available and can throw an exception or return None. This change will handle both jaxlib.xla_extension.XlaRuntimeError if the device is not a PjRt addressable device or KeyError if the memstats returns None if they are not available.
Memstats are not guaranteed to be available and can throw an exception or return None. This change will handle both
jaxlib.xla_extension.XlaRuntimeErrorif the device is not a PjRt addressable device orKeyErrorif the memstats returns None if they are not available.