in bindings/python/src/lib.rs [325:351]
fn extract_bound(ob: &PyBound<'source, PyAny>) -> PyResult<Self> {
if let Ok(name) = ob.extract::<String>() {
match name.as_str() {
"cpu" => Ok(Device::Cpu),
"cuda" => Ok(Device::Cuda(0)),
"mps" => Ok(Device::Mps),
"npu" => Ok(Device::Npu(0)),
"xpu" => Ok(Device::Xpu(0)),
"xla" => Ok(Device::Xla(0)),
"mlu" => Ok(Device::Mlu(0)),
"hpu" => Ok(Device::Hpu(0)),
name if name.starts_with("cuda:") => parse_device(name).map(Device::Cuda),
name if name.starts_with("npu:") => parse_device(name).map(Device::Npu),
name if name.starts_with("xpu:") => parse_device(name).map(Device::Xpu),
name if name.starts_with("xla:") => parse_device(name).map(Device::Xla),
name if name.starts_with("mlu:") => parse_device(name).map(Device::Mlu),
name if name.starts_with("hpu:") => parse_device(name).map(Device::Hpu),
name => Err(SafetensorError::new_err(format!(
"device {name} is invalid"
))),
}
} else if let Ok(number) = ob.extract::<usize>() {
Ok(Device::Anonymous(number))
} else {
Err(SafetensorError::new_err(format!("device {ob} is invalid")))
}
}