

from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
  import torch

A__										=      logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
  from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):





  def 				__a							(       cls						):
      return cls





def 				__a							(       __A			,							__A = None			,							__A = None			,							__A = None			,							__A = None			,							):
    SCREAMING_SNAKE_CASE__         =    device
    SCREAMING_SNAKE_CASE__         =    shape[0]

    SCREAMING_SNAKE_CASE__         =    layout or torch.strided
    SCREAMING_SNAKE_CASE__         =    device or torch.device("cpu"						)

    if generator is not None:
        SCREAMING_SNAKE_CASE__         =    generator.device.type if not isinstance(__A			,							__A						) else generator[0].device.type
        if gen_device_type != device.type and gen_device_type == "cpu":
            SCREAMING_SNAKE_CASE__         =    "cpu"
            if device != "mps":
                logger.info(
                    f"The passed generator was created on 'cpu' even though a tensor on {device} was expected."
                    f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"
                    f" slighly speed up this function by passing a generator that was created on the {device} device."						)
        elif gen_device_type != device.type and gen_device_type == "cuda":
            raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}."						)

    if isinstance(__A			,							__A						):
        SCREAMING_SNAKE_CASE__         =    (1,) + shape[1:]
        SCREAMING_SNAKE_CASE__         =    [
            torch.randn(__A			,							generator=generator[i]			,							device=__A			,							dtype=__A			,							layout=__A						)
            for i in range(__A						)
        ]
        SCREAMING_SNAKE_CASE__         =    torch.cat(__A			,							dim=0						).to(__A						)
    else:
        SCREAMING_SNAKE_CASE__         =    torch.randn(__A			,							generator=__A			,							device=__A			,							dtype=__A			,							layout=__A						).to(__A						)

    return latents





def 				__a							(       __A						):
    if is_torch_version("<"			,							"2.0.0"						) or not hasattr(__A			,							"_dynamo"						):
        return False
    return isinstance(__A			,							torch._dynamo.eval_frame.OptimizedModule						)





from typing import Optional

from ..utils.logging import get_logger
from .audio_classificiation import AudioClassification
from .automatic_speech_recognition import AutomaticSpeechRecognition
from .base import TaskTemplate
from .image_classification import ImageClassification
from .language_modeling import LanguageModeling
from .question_answering import QuestionAnsweringExtractive
from .summarization import Summarization
from .text_classification import TextClassification


A__										=      [
    "AutomaticSpeechRecognition",
    "AudioClassification",
    "ImageClassification",
    "LanguageModeling",
    "QuestionAnsweringExtractive",
    "Summarization",
    "TaskTemplate",
    "TextClassification",
]

A__										=      get_logger(__name__)


A__										=      {
    AutomaticSpeechRecognition.task: AutomaticSpeechRecognition,
    AudioClassification.task: AudioClassification,
    ImageClassification.task: ImageClassification,
    LanguageModeling.task: LanguageModeling,
    QuestionAnsweringExtractive.task: QuestionAnsweringExtractive,
    Summarization.task: Summarization,
    TextClassification.task: TextClassification,
}





def 				__a							(       __A						):
    SCREAMING_SNAKE_CASE__         =    task_template_dict.get("task"						)
    if task_name is None:
        logger.warning(f"Couldn't find template for task '{task_name}'. Available templates: {list(__A						)}"						)
        return None
    SCREAMING_SNAKE_CASE__         =    NAMEaTEMPLATE.get(__A						)
    return template.from_dict(__A						)


