Move PJRT Python APIs out of torch_xla.experimental.*

[will-cromar]

Reorganize the experimental PJRT Python APIs.

• Add new _internal module for APIs that are well-tested, but likely to change. I moved device-specific logic here, since I expect to rework it in the near future. All of these functions are mainly used for framework development. In general, users shouldn't have to call them directly.
• Most functionality that interacts directly with the runtime is moved into the new torch.runtime module.
• Added deprecation module to register deprecated aliases for all public functions that are moving out into other parts of of torch_xla.
• Print a warning once per deprecated function. For example:

>>> from torch_xla.experimental import tpu
>>> tpu.num_available_chips()
WARNING:root:torch_xla.experimental.tpu.num_available_chips is deprecated. Use torch_xla._internal.tpu.num_available_chips instead.
4
>>> tpu.num_available_chips()
4
>>> tpu.version()
WARNING:root:torch_xla.experimental.tpu.version is deprecated. Use torch_xla._internal.tpu.version instead.
4

• Update references across the repository.

Summary of new modules:

• torch_xla.runtime
• torch_xla._internal.tpu
• torch_xla._internal.gpu
• torch_xla._internal.pjrt

[JackCaoG]

should we renamed these experimental files?

[will-cromar]

Yes, good catch.

[JackCaoG]

I think we have similar functions in torch_xla.core.xla_model, do we want to do some clean up?

[JackCaoG]

I want to take this chance to do some clean up. I am always confuse what function to call for local ordinal, gloabal ordinal, worla_size etc and what do they really mean in a pod context. If we can restructure those api a bit and maybe have those apis in this runtime instead that would be nice...(random idea, might need more thinking)

[JackCaoG]

xla_model becomes a kitch sink and we put random things in it, if we can move all runtime related bits in this file it is actually nicer..

[will-cromar]

Discussed this offline. We'll start to move APIs that interact directly with the runtime to a new module, and leave any modeling-related APIs in xla_model. I moved the PJRT version of rendezvous back to xla_model, and the old rendezvous is will be an alias of that implementation when PJRT is enabled.

[JackCaoG]

do we want to change preview to stable here?

[will-cromar]

Good catch. Done.

[JackCaoG]

LGTM, but prefer to merge on Monday