Expose default random number generator for XLA devices

[lqf96]

🚀 Feature

It would be nice to expose default random generator for XLA devices, just like CPU (at torch.default_generator) and GPU (at torch.cuda.default_generators) devices in PyTorch. The default generators can live in e.g. torch_xla.core.random.default_generators and the torch_xla.core.random can additionally provide a number of methods to get, set and seed the random number generator states (e.g. seed, manual_seed, get_rng_state, set_rng_state).

Motivation

This will be helpful for reproducing the results of neural networks when using TPUs.

[JackCaoG]

You can set the seed via

xla/torch_xla/core/xla_model.py

Line 917 in df0a277

def set_rng_state(seed, device=None):

some example in

xla/test/test_operations.py

Line 81 in 132481b

def _set_rng_seed(seed):

Will this be enough for your use case?

[lqf96]

@JackCaoG Yeah I know that... But I have some code that operates on torch.Generators, and it would be more convenient if default_generators are provided and torch.random.manual_seed extends support to XLA devices.

[JackCaoG]

Pytorch/XLA runs on XLA device, we used RngBitGenerator OP to generate the random number (where we can only select an algorithm and set the initial state/seed). We don't use the torch.Generator and I can't think of a way to fit this into the Pytorch/XLA.

Even if we can call generator on CPU and then feed the constant to the graph, the constant will change every step and we will keep recompiling the graph.

[lqf96]

Ok I read some PyTorch and XLA source code and noticed that this will be, in principle, doable. In PyTorch each torch.Generator is a aten::Generator in C++, which in turn is implemented by a device-specific subclass of c10::GeneratorImpl. Also notice that in XLA, all random operations return both results and new random states, as shown by the RngOutput struct, but as of now PyTorch XLA ignores the random states. So one idea to implement this will be:

1. Make a subclass of c10::GeneratorImpl for XLA devices, possibly called XLAGeneratorImpl and initialize this generator with a XLA tensor as the RNG state which internally contains a XLA graph node.
2. When performing random-related operations, use the current XLAGeneratorImpl state as the input RNG state node and then replace the XLAGeneratorImpl state with the output RNG state node.
3. Under this design the XLA Generator will also be lazy and the actual random state is only revealed (or executed) when APIs like aten::Generator::current_seed or Generator.get_state is called.

I can look into this and try to come up with a PR if I have time.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.