def list_devices(self, device_type: str or None = None) -> List[ComputeDevice]:
        devices = []
        for jax_dev in jax.devices():
            jax_dev_type = jax_dev.platform.upper()
            if device_type is None or device_type == jax_dev_type:
                description = f"id={jax_dev.id}"
                devices.append(ComputeDevice(self, jax_dev.device_kind, jax_dev_type, -1, -1, description, jax_dev))
        return devices