from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )     ->   str:
    for param, grad_param in zip(model_a.parameters()    ,  model_b.parameters()    ):
        if not param.requires_grad:
            continue
        if not did_step:
            # Grads should not be in sync
            assert (
                torch.allclose(param.grad    ,  grad_param.grad    ) is False
            ), f"""Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"""
        else:
            # Grads should be in sync
            assert (
                torch.allclose(param.grad    ,  grad_param.grad    ) is True
            ), f"""Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"""






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__=True    )     ->   Union[str, Any]:
    model.train()
    __lowerCamelCase          =  model(UpperCamelCase__    )
    __lowerCamelCase          =  F.mse_loss(UpperCamelCase__    ,  target.to(output.device    )    )
    if not do_backward:
        loss /= accelerator.gradient_accumulation_steps
        loss.backward()
    else:
        accelerator.backward(UpperCamelCase__    )






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__=False    )     ->   Optional[Any]:
    set_seed(42    )
    __lowerCamelCase          =  RegressionModel()
    __lowerCamelCase          =  deepcopy(UpperCamelCase__    )
    __lowerCamelCase          =  RegressionDataset(length=80    )
    __lowerCamelCase          =  DataLoader(UpperCamelCase__    ,  batch_size=16    )
    model.to(accelerator.device    )
    if sched:
        __lowerCamelCase          =  AdamW(params=model.parameters()    ,  lr=1E-3    )
        __lowerCamelCase          =  AdamW(params=ddp_model.parameters()    ,  lr=1E-3    )
        __lowerCamelCase          =  LambdaLR(UpperCamelCase__    ,  lr_lambda=lambda UpperCamelCase__    : epoch**0.6_5    )
        __lowerCamelCase          =  LambdaLR(UpperCamelCase__    ,  lr_lambda=lambda UpperCamelCase__    : epoch**0.6_5    )
    # Make a copy of `model`
    if sched:
        __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.prepare(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
    else:
        __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.prepare(UpperCamelCase__    ,  UpperCamelCase__    )
    if sched:
        return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
    return model, ddp_model, dataloader






def __lowerCAmelCase			(							UpperCamelCase__    )     ->   Tuple:
    # Test when on a single CPU or GPU that the context manager does nothing
    __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  get_training_setup(UpperCamelCase__    )
    # Use a single batch
    __lowerCamelCase     ,  __lowerCamelCase          =  next(iter(UpperCamelCase__    )    ).values()
    for iteration in range(3    ):
        # Gather the distributed inputs and targs for the base model
        __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.gather((ddp_input, ddp_target)    )
        __lowerCamelCase     ,  __lowerCamelCase          =  input.to(accelerator.device    ), target.to(accelerator.device    )
        # Perform our initial ground truth step in non "DDP"
        step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        # Do "gradient accumulation" (noop)
        if iteration % 2 == 0:
            # Accumulate grads locally
            with accelerator.no_sync(UpperCamelCase__    ):
                step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        else:
            # Sync grads
            step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )

        # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
        check_model_parameters(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        for param, ddp_param in zip(model.parameters()    ,  ddp_model.parameters()    ):
            if not param.requires_grad:
                continue
            assert torch.allclose(
                param.grad    ,  ddp_param.grad    ), f"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
        torch.manual_seed(13_37 + iteration    )
        __lowerCamelCase          =  ddp_input[torch.randperm(len(UpperCamelCase__    )    )]






def __lowerCAmelCase			(							UpperCamelCase__    )     ->   int:
    # Test on distributed setup that context manager behaves properly
    __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  get_training_setup(UpperCamelCase__    )
    # Use a single batch
    __lowerCamelCase     ,  __lowerCamelCase          =  next(iter(UpperCamelCase__    )    ).values()
    for iteration in range(3    ):
        # Gather the distributed inputs and targs for the base model
        __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.gather((ddp_input, ddp_target)    )
        __lowerCamelCase     ,  __lowerCamelCase          =  input.to(accelerator.device    ), target.to(accelerator.device    )
        # Perform our initial ground truth step in non "DDP"
        step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        # Do "gradient accumulation" (noop)
        if iteration % 2 == 0:
            # Accumulate grads locally
            with accelerator.no_sync(UpperCamelCase__    ):
                step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        else:
            # Sync grads
            step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )

        # DDP model and model should only be in sync when not (iteration % 2 == 0)
        for param, ddp_param in zip(model.parameters()    ,  ddp_model.parameters()    ):
            if not param.requires_grad:
                continue
            if iteration % 2 == 0:
                # Grads should not be in sync
                assert (
                    torch.allclose(param.grad    ,  ddp_param.grad    ) is False
                ), f"""Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""
            else:
                # Grads should be in sync
                assert (
                    torch.allclose(param.grad    ,  ddp_param.grad    ) is True
                ), f"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
        torch.manual_seed(13_37 + iteration    )
        __lowerCamelCase          =  ddp_input[torch.randperm(len(UpperCamelCase__    )    )]






def __lowerCAmelCase			(							UpperCamelCase__=False    ,  UpperCamelCase__=False    )     ->   int:
    __lowerCamelCase          =  Accelerator(
        split_batches=UpperCamelCase__    ,  dispatch_batches=UpperCamelCase__    ,  gradient_accumulation_steps=2    )
    # Test that context manager behaves properly
    __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  get_training_setup(UpperCamelCase__    )
    for iteration, batch in enumerate(UpperCamelCase__    ):
        __lowerCamelCase     ,  __lowerCamelCase          =  batch.values()
        # Gather the distributed inputs and targs for the base model
        __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.gather((ddp_input, ddp_target)    )
        __lowerCamelCase     ,  __lowerCamelCase          =  input.to(accelerator.device    ), target.to(accelerator.device    )
        # Perform our initial ground truth step in non "DDP"
        step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        # Do "gradient accumulation" (noop)
        with accelerator.accumulate(UpperCamelCase__    ):
            step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )

        # DDP model and model should only be in sync when not (iteration % 2 == 0)
        for param, ddp_param in zip(model.parameters()    ,  ddp_model.parameters()    ):
            if not param.requires_grad:
                continue
            if ((iteration + 1) % 2 == 0) or (iteration == len(UpperCamelCase__    ) - 1):
                # Grads should be in sync
                assert (
                    torch.allclose(param.grad    ,  ddp_param.grad    ) is True
                ), f"""Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""
            else:
                # Grads should not be in sync
                assert (
                    torch.allclose(param.grad    ,  ddp_param.grad    ) is False
                ), f"""Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
        torch.manual_seed(13_37 + iteration    )
        __lowerCamelCase          =  ddp_input[torch.randperm(len(UpperCamelCase__    )    )]
    GradientState._reset_state()






def __lowerCAmelCase			(							UpperCamelCase__=False    ,  UpperCamelCase__=False    )     ->   Tuple:
    __lowerCamelCase          =  Accelerator(
        split_batches=UpperCamelCase__    ,  dispatch_batches=UpperCamelCase__    ,  gradient_accumulation_steps=2    )
    # Test that context manager behaves properly
    __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  get_training_setup(UpperCamelCase__    ,  UpperCamelCase__    )
    for iteration, batch in enumerate(UpperCamelCase__    ):
        __lowerCamelCase     ,  __lowerCamelCase          =  batch.values()
        # Gather the distributed inputs and targs for the base model
        __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.gather((ddp_input, ddp_target)    )
        __lowerCamelCase     ,  __lowerCamelCase          =  input.to(accelerator.device    ), target.to(accelerator.device    )
        # Perform our initial ground truth step in non "DDP"
        model.train()
        ddp_model.train()
        step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        opt.step()

        if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(UpperCamelCase__    )):
            if split_batches:
                sched.step()
            else:
                for _ in range(accelerator.num_processes    ):
                    sched.step()
        opt.zero_grad()
        # Perform gradient accumulation under wrapper
        with accelerator.accumulate(UpperCamelCase__    ):
            step_model(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
            ddp_opt.step()
            ddp_sched.step()
            ddp_opt.zero_grad()

        # Learning rates should be the same
        assert (
            opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
        ), f"""Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n"""
        __lowerCamelCase          =  (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(UpperCamelCase__    ))
        if accelerator.num_processes > 1:
            check_model_parameters(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        # Shuffle ddp_input on each iteration
        torch.manual_seed(13_37 + iteration    )
    GradientState._reset_state()






def __lowerCAmelCase			(							)     ->   List[str]:
    __lowerCamelCase          =  Accelerator()

    __lowerCamelCase          =  RegressionDataset(length=80    )
    __lowerCamelCase          =  DataLoader(UpperCamelCase__    ,  batch_size=16    )
    __lowerCamelCase          =  RegressionDataset(length=96    )
    __lowerCamelCase          =  DataLoader(UpperCamelCase__    ,  batch_size=16    )
    __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.prepare(UpperCamelCase__    ,  UpperCamelCase__    )
    assert accelerator.gradient_state.active_dataloader is None
    for iteration, _ in enumerate(UpperCamelCase__    ):
        assert id(accelerator.gradient_state.active_dataloader    ) == id(UpperCamelCase__    )
        if iteration < len(UpperCamelCase__    ) - 1:
            assert not accelerator.gradient_state.end_of_dataloader
            if iteration == 1:
                for batch_num, _ in enumerate(UpperCamelCase__    ):
                    assert id(accelerator.gradient_state.active_dataloader    ) == id(UpperCamelCase__    )
                    if batch_num < len(UpperCamelCase__    ) - 1:
                        assert not accelerator.gradient_state.end_of_dataloader
                    else:
                        assert accelerator.gradient_state.end_of_dataloader
        else:
            assert accelerator.gradient_state.end_of_dataloader
    assert accelerator.gradient_state.active_dataloader is None






def __lowerCAmelCase			(							)     ->   Optional[int]:
    __lowerCamelCase          =  Accelerator()
    __lowerCamelCase          =  accelerator.state
    if state.local_process_index == 0:
        print('''**Test `accumulate` gradient accumulation with dataloader break**'''    )
    test_dataloader_break()
    if state.distributed_type == DistributedType.NO:
        if state.local_process_index == 0:
            print('''**Test NOOP `no_sync` context manager**'''    )
        test_noop_sync(UpperCamelCase__    )
    if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
        if state.local_process_index == 0:
            print('''**Test Distributed `no_sync` context manager**'''    )
        test_distributed_sync(UpperCamelCase__    )
    if state.distributed_type == DistributedType.MULTI_GPU:
        for split_batch in [True, False]:
            for dispatch_batches in [True, False]:
                if state.local_process_index == 0:
                    print(
                        '''**Test `accumulate` gradient accumulation, '''    ,  f"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""    ,  )
                test_gradient_accumulation(UpperCamelCase__    ,  UpperCamelCase__    )

    # Currently will break on torch 2.0 +, need to investigate why
    if is_torch_version('''<'''    ,  '''2.0'''    ) or state.distributed_type == DistributedType.NO:
        if state.local_process_index == 0:
            print(
                '''**Test `accumulate` gradient accumulation with optimizer and scheduler, '''    ,  '''`split_batches=False`, `dispatch_batches=False`**'''    ,  )
        test_gradient_accumulation_with_opt_and_scheduler()
        if state.distributed_type == DistributedType.MULTI_GPU:
            for split_batch in [True, False]:
                for dispatch_batches in [True, False]:
                    if not split_batch and not dispatch_batches:
                        continue
                    if state.local_process_index == 0:
                        print(
                            '''**Test `accumulate` gradient accumulation with optimizer and scheduler, '''    ,  f"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""    ,  )
                    test_gradient_accumulation_with_opt_and_scheduler(UpperCamelCase__    ,  UpperCamelCase__    )






def __lowerCAmelCase			(							UpperCamelCase__    )     ->   List[str]:
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
  main()

'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase												={
    "configuration_whisper": ["WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP", "WhisperConfig", "WhisperOnnxConfig"],
    "feature_extraction_whisper": ["WhisperFeatureExtractor"],
    "processing_whisper": ["WhisperProcessor"],
    "tokenization_whisper": ["WhisperTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=["WhisperTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "WhisperForConditionalGeneration",
      "WhisperModel",
      "WhisperPreTrainedModel",
      "WhisperForAudioClassification",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFWhisperForConditionalGeneration",
      "TFWhisperModel",
      "TFWhisperPreTrainedModel",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "FlaxWhisperForConditionalGeneration",
      "FlaxWhisperModel",
      "FlaxWhisperPreTrainedModel",
      "FlaxWhisperForAudioClassification",
  ]


if TYPE_CHECKING:
  from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
  from .feature_extraction_whisper import WhisperFeatureExtractor
  from .processing_whisper import WhisperProcessor
  from .tokenization_whisper import WhisperTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_whisper_fast import WhisperTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_whisper import (
        WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
        WhisperForAudioClassification,
        WhisperForConditionalGeneration,
        WhisperModel,
        WhisperPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_whisper import (
        TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFWhisperForConditionalGeneration,
        TFWhisperModel,
        TFWhisperPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_whisper import (
        FlaxWhisperForAudioClassification,
        FlaxWhisperForConditionalGeneration,
        FlaxWhisperModel,
        FlaxWhisperPreTrainedModel,
    )

else:
  import sys

  __UpperCAmelCase												=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
