Autocast Policy for new xla backend.

[avizon-aws]

🚀 Feature

I want to add support for a new autocast policy for the NEURON backend.

Motivation

In my usecase, the device is an XLA device but the backend is different. Currently the autocast policy being used is : https://github.com/pytorch/xla/blob/master/torch_xla/csrc/autocast_mode.cpp, but if I want to have a different policy for TPU and NEURON, then I am not sure if it is currently possible. It would be nice to be able to support this feature, because even though both are xla devices, each backend can maintain its own set of operations in a different precision as needed.

Pitch

I want to be able to create a new autocast policy for my backend of interest. I also want to be able to inherit the policy from the xla device, but want the feature to be able to override it as per the need.

Alternatives

Additional context

[ysiraichi]

Thank you for your suggestion.

Although, there's probably no easy way to overwrite the autocast policy, separately, for each XLA backend, I believe you can make use of the custom dispatch key PrivateUse1. Specifically, you could:

• Register autocast kernels for AutocastPrivateUse1
• Create a new TensorImpl that inherits from XLATensorImpl
• Replace their dispatch key from AutocastXLA to AutocastPrivateUse1

[miladm]

@avizon-aws: following up - are you unblocked here?

[avizon-aws]

@miladm , i am unblocked here, this was more for future purposes if the need arises.