def derive_rand(rand: th.Generator, device: Device) -> th.Generator:
device = as_device(device)
# Return existing random number generator for the same device
if rand.device == device:
return rand
# Create and seed new random number generator
rand_new = th.Generator(device)
rand_new.manual_seed(rand.seed() & _SEED_MASK)
return rand_new
def use_rand(rand: th.Generator, **kwargs: Any):
# Fork and seed current global random state
with th.random.fork_rng(devices=(rand.device, ), **kwargs):
th.random.manual_seed(rand.seed() & _SEED_MASK)
yield