A helper function to create random tensors on the desired `device` with the desired `dtype`. When passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor is always created on the CPU.
(
logger,
shape: Union[Tuple, List],
generator: Optional[Union[List['torch.Generator'],
'torch.Generator']] = None,
device: Optional['torch.device'] = None,
dtype: Optional['torch.dtype'] = None,
layout: Optional['torch.layout'] = None,
)
| 6 | |
| 7 | |
| 8 | def randn_tensor( |
| 9 | logger, |
| 10 | shape: Union[Tuple, List], |
| 11 | generator: Optional[Union[List['torch.Generator'], |
| 12 | 'torch.Generator']] = None, |
| 13 | device: Optional['torch.device'] = None, |
| 14 | dtype: Optional['torch.dtype'] = None, |
| 15 | layout: Optional['torch.layout'] = None, |
| 16 | ): |
| 17 | """A helper function to create random tensors on the desired `device` with |
| 18 | the desired `dtype`. |
| 19 | |
| 20 | When passing a list of generators, you can seed each batch size |
| 21 | individually. If CPU generators are passed, the tensor is always created on |
| 22 | the CPU. |
| 23 | """ |
| 24 | # device on which tensor is created defaults to device |
| 25 | rand_device = device |
| 26 | batch_size = shape[0] |
| 27 | |
| 28 | layout = layout or torch.strided |
| 29 | device = device or torch.device('cpu') |
| 30 | |
| 31 | if generator is not None: |
| 32 | gen_device_type = generator.device.type if not isinstance( |
| 33 | generator, list) else generator[0].device.type |
| 34 | if gen_device_type != device.type and gen_device_type == 'cpu': |
| 35 | rand_device = 'cpu' |
| 36 | if device != 'mps': |
| 37 | logger.info( |
| 38 | f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." |
| 39 | f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" |
| 40 | f' slightly speed up this function by passing a generator that was created on the {device} device.' |
| 41 | ) |
| 42 | elif gen_device_type != device.type and gen_device_type == 'cuda': |
| 43 | raise ValueError( |
| 44 | f'Cannot generate a {device} tensor from a generator of type {gen_device_type}.' |
| 45 | ) |
| 46 | |
| 47 | # make sure generator list of length 1 is treated like a non-list |
| 48 | if isinstance(generator, list) and len(generator) == 1: |
| 49 | generator = generator[0] |
| 50 | |
| 51 | if isinstance(generator, list): |
| 52 | shape = (1, ) + shape[1:] |
| 53 | latents = [ |
| 54 | torch.randn( |
| 55 | shape, |
| 56 | generator=generator[i], |
| 57 | device=rand_device, |
| 58 | dtype=dtype, |
| 59 | layout=layout) for i in range(batch_size) |
| 60 | ] |
| 61 | latents = torch.cat(latents, dim=0).to(device) |
| 62 | else: |
| 63 | latents = torch.randn( |
| 64 | shape, |
| 65 | generator=generator, |