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.
(
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,
)
| 36 | |
| 37 | |
| 38 | def randn_tensor( |
| 39 | shape: Union[Tuple, List], |
| 40 | generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None, |
| 41 | device: Optional["torch.device"] = None, |
| 42 | dtype: Optional["torch.dtype"] = None, |
| 43 | layout: Optional["torch.layout"] = None, |
| 44 | ): |
| 45 | """A helper function to create random tensors on the desired `device` with the desired `dtype`. When |
| 46 | passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor |
| 47 | is always created on the CPU. |
| 48 | """ |
| 49 | # device on which tensor is created defaults to device |
| 50 | rand_device = device |
| 51 | batch_size = shape[0] |
| 52 | |
| 53 | layout = layout or torch.strided |
| 54 | device = device or torch.device("cpu") |
| 55 | |
| 56 | if generator is not None: |
| 57 | gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type |
| 58 | if gen_device_type != device.type and gen_device_type == "cpu": |
| 59 | rand_device = "cpu" |
| 60 | if device != "mps": |
| 61 | logger.info( |
| 62 | f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." |
| 63 | f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" |
| 64 | f" slighly speed up this function by passing a generator that was created on the {device} device." |
| 65 | ) |
| 66 | elif gen_device_type != device.type and gen_device_type == "cuda": |
| 67 | raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.") |
| 68 | |
| 69 | # make sure generator list of length 1 is treated like a non-list |
| 70 | if isinstance(generator, list) and len(generator) == 1: |
| 71 | generator = generator[0] |
| 72 | |
| 73 | if isinstance(generator, list): |
| 74 | shape = (1,) + shape[1:] |
| 75 | latents = [ |
| 76 | torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout) |
| 77 | for i in range(batch_size) |
| 78 | ] |
| 79 | latents = torch.cat(latents, dim=0).to(device) |
| 80 | else: |
| 81 | latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device) |
| 82 | |
| 83 | return latents |
| 84 | |
| 85 | |
| 86 | def is_compiled_module(module) -> bool: |