declare-lab/tango2

1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""15PyTorch utilities: Utilities related to PyTorch16"""17from typing import List, Optional, Tuple, Union18 19from . import logging20from .import_utils import is_torch_available, is_torch_version21 22 23if is_torch_available():24    import torch25 26logger = logging.get_logger(__name__)  # pylint: disable=invalid-name27 28 29def randn_tensor(30    shape: Union[Tuple, List],31    generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None,32    device: Optional["torch.device"] = None,33    dtype: Optional["torch.dtype"] = None,34    layout: Optional["torch.layout"] = None,35):36    """This is a helper function that allows to create random tensors on the desired `device` with the desired `dtype`. When37    passing a list of generators one can seed each batched size individually. If CPU generators are passed the tensor38    will always be created on CPU.39    """40    # device on which tensor is created defaults to device41    rand_device = device42    batch_size = shape[0]43 44    layout = layout or torch.strided45    device = device or torch.device("cpu")46 47    if generator is not None:48        gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type49        if gen_device_type != device.type and gen_device_type == "cpu":50            rand_device = "cpu"51            if device != "mps":52                logger.info(53                    f"The passed generator was created on 'cpu' even though a tensor on {device} was expected."54                    f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"55                    f" slighly speed up this function by passing a generator that was created on the {device} device."56                )57        elif gen_device_type != device.type and gen_device_type == "cuda":58            raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.")59 60    if isinstance(generator, list):61        shape = (1,) + shape[1:]62        latents = [63            torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout)64            for i in range(batch_size)65        ]66        latents = torch.cat(latents, dim=0).to(device)67    else:68        latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device)69 70    return latents71 72 73def is_compiled_module(module):74    """Check whether the module was compiled with torch.compile()"""75    if is_torch_version("<", "2.0.0") or not hasattr(torch, "_dynamo"):76        return False77    return isinstance(module, torch._dynamo.eval_frame.OptimizedModule)78