keras/src/backend/jax/distribution_lib.py/_to_jax_device
def _to_jax_device(device_name):
    if isinstance(device_name, jax.Device):
        return device_name
    device_type, device_id = device_name.split(":")

    devices = jax.devices(backend=device_type)
    for device in devices:
        if device.platform == device_type and device.id == int(device_id):
            return device
    raise ValueError(f"Device not found: {device_name}")