Support for XLA devices

[guoquan]

Description

XLA is an abstraction layer of computation graph for better efficiency, consistency, portability, and a lot as they claim. However, the most significant feature is to enable access to google's TPU or other customized accelerator hardware that uses the same abstraction.

Sample code to use XLA device could be:

from mxnet import nd
from mxnet_xla import xla
x = nd.ones((4,5), ctx=xla.xla_device())
print(x)

References

• https://www.tensorflow.org/xla
• https://github.com/pytorch/xla/blob/master/API_GUIDE.md
• https://mxnet.apache.org/api/python/docs/tutorials/getting-started/crash-course/6-use_gpus.html

[pengzhao-intel]

Good idea! We're evaluating the possibility of XLA recently :)

[cjolivier01]

As a start, it would be good to simply be able to generate an HloModuleProto protobuf file.