Reset PRNG key placement in distributed_array tutorial

[wdhongtw]

jax_jit stores its output key in env.prng_key, committing it to the last call's devices, so switching between sharded and single-device inputs fails with "incompatible devices". Reset the key to an uncommitted copy at each switch, and note the limitation in interop.

---

The error from https://github.com/google/torchax/actions/workflows/static.yml recently, is that the PRNG env.prng_key is now committed after first use. If the first use is computation on sharded device, and then the second use is on some single device (or vise versa), the committed PRNG cause runtime error.

See also: #95 for more background where this behavior was introduced.

import jax
jax.config.update("jax_num_cpu_devices", 4)  # simulate 4 devices on CPU

import torch
import torchax
from jax.sharding import NamedSharding, PartitionSpec as P

torchax.enable_globally()

mesh = jax.make_mesh((4,), ("x",))
x = torch.ones(4, 4, device="jax")
sharded = x.apply_jax(jax.device_put, NamedSharding(mesh, P("x")))
single = x.apply_jax(jax.device_put, jax.devices()[0])

fn = torchax.interop.jax_jit(lambda a: (a * 2).sum())
fn(sharded)  # env.prng_key is now tied to all 4 devices
fn(single)   # ValueError: Received incompatible devices for jitted computation.
             #   ... device ids [0] ... key<fry>[] ... device ids [0, 1, 2, 3]

I have try some solution in https://github.com/google/torchax/tree/prng-split-before-call, but I think that takes some time to get a mature and solid solution. Let's treat current implementation/design as a limitation and go back later.