Fix XPU backend availability by importing IPEX

[codegen-sh]

Problem

When running kmeans_aces_xpu.slurm with XPU devices, torch.xpu.is_available() returns False even though XPU devices are properly allocated by SLURM:

DEBUG | exp.activations:load_activations_and_init_dist:813 - Backend: <module 'torch.xpu' from '...'>
DEBUG | exp.activations:load_activations_and_init_dist:814 - Backend is available: False
...
AssertionError: CPU-only not supported yet :( Device xpu not available.

Root Cause

Intel Extension for PyTorch (IPEX) must be imported before any XPU operations to register the XPU backend with PyTorch's device system. Without this import:

• torch.xpu module exists but is not fully initialized
• torch.xpu.is_available() returns False
• XPU device creation and operations fail

This is a requirement of Intel's XPU backend that's not obvious from the documentation.

Solution

Added lazy IPEX import to core/device.py:

1. _ensure_ipex_imported(): New helper function that imports IPEX on first use

   • Uses module-level flag to ensure import happens only once
   • Gracefully handles missing IPEX (will fail later with clear error)

2. Modified get_backend(): Calls _ensure_ipex_imported() before returning th.xpu

   • Ensures backend is registered before any operations

3. Modified get_device(): Calls _ensure_ipex_imported() for XPU devices

   • Ensures backend is registered before device creation

This fix ensures the XPU backend is properly initialized whenever XPU functionality is accessed, making torch.xpu.is_available() return True when XPU devices are present.

Testing

All checks passed:

• uv run ruff check . - No linting errors
• uv run ruff format . - Code formatted correctly

References

• Intel Extension for PyTorch is required for XPU support
• The import registers the XPU device type with PyTorch's backend system

[codegen-sh]

🔍 Broken test auto-fixer • Learn more

Check Suite	Agent	Status	Commit	Time
GitHub Actions	Agent	Fix ✅	ab39f13	Oct 30, 02:59:25 UTC

💻 View my work • 🛑 Stop • 🚫 Ban all checks

⚙️ Check suite settings