xla icon indicating copy to clipboard operation
xla copied to clipboard

Add XLAHooksInterface to PyTorch for XLA

Open iwknow opened this issue 3 months ago • 0 comments

🚀 Feature

Tytorch core uses AcceleratorHooksInterface to interact with different accelerators. For example, different accelerator has their own device-specific ways to generate random number. Instead of adding device-specific code to Pytorch core, it uses method like getDefaultGenerator from AcceleratorHooksInterface to get a random number generator. different accelerator should implement this interface so that torch can do the right thing on the device. Currently, most of the accelerators have their own implementation of AcceleratorHooksInterface. Their implementations can be found here

Motivation

This feature is directly needed by enabling generator on XLA device (details). More importantly, once created, it creates an avenue for pytorch to directly interact with XLA device.

Pitch

A PR has already been created for this change and approved by Pytorch folks. We need someone from XLA side to review: https://github.com/pytorch/pytorch/pull/161369

iwknow avatar Oct 02 '25 16:10 iwknow