🚀 The feature, motivation and pitch

Motivation

Distributed tests currently contain many hardcoded device count lookups, for example:

torch.cuda.device_count()
torch.hpu.device_count()
torch.xpu.device_count()
torch.accelerator.device_count()

These lookups are scattered across distributed test bases, decorators, helpers, and world-size setup logic.

This makes distributed test infrastructure harder to reuse for out-of-tree backends and harder to make device-generic.

torch.cuda.device_count(), torch.hpu.device_count(), and torch.xpu.device_count() hardcode specific device families, while torch.accelerator.device_count() is tied to the current accelerator rather than the instantiated device-specific test class's device_type.

Device count is fundamentally a device-type-level property. For instantiated device-specific test classes, device count should instead be derived from the test class's device_type.

Proposal

Add a device_count() hook to DeviceTypeTestBase:

class DeviceTypeTestBase(TestCase):
    @classmethod
    def device_count(cls) -> int:
        return torch.get_device_module(cls.device_type).device_count()

This provides a single device-type-aware entry point for distributed test infrastructure to query device count.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @fffrog

Alternatives

No response

Additional context

No response