"""simple docstring"""


import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        FlaubertForMultipleChoice,
        FlaubertForQuestionAnswering,
        FlaubertForQuestionAnsweringSimple,
        FlaubertForSequenceClassification,
        FlaubertForTokenClassification,
        FlaubertModel,
        FlaubertWithLMHeadModel,
    )
    from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST


class 				a (						lowerCAmelCase_      ):


    def __init__(    self		:     List[str]      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     int=13      , __lowerCAmelCase		:     Optional[Any]=7      , __lowerCAmelCase		:     Union[str, Any]=True      , __lowerCAmelCase		:     str=True      , __lowerCAmelCase		:     Optional[int]=True      , __lowerCAmelCase		:     Tuple=True      , __lowerCAmelCase		:     Union[str, Any]=True      , __lowerCAmelCase		:     str=False      , __lowerCAmelCase		:     Optional[Any]=False      , __lowerCAmelCase		:     List[str]=False      , __lowerCAmelCase		:     List[Any]=2      , __lowerCAmelCase		:     Optional[int]=99      , __lowerCAmelCase		:     str=0      , __lowerCAmelCase		:     Any=32      , __lowerCAmelCase		:     List[Any]=5      , __lowerCAmelCase		:     Dict=4      , __lowerCAmelCase		:     str=0.1      , __lowerCAmelCase		:     Optional[int]=0.1      , __lowerCAmelCase		:     Union[str, Any]=512      , __lowerCAmelCase		:     str=12      , __lowerCAmelCase		:     Optional[Any]=2      , __lowerCAmelCase		:     List[Any]=0.02      , __lowerCAmelCase		:     Optional[int]=3      , __lowerCAmelCase		:     Any=4      , __lowerCAmelCase		:     int="last"      , __lowerCAmelCase		:     Dict=None      , __lowerCAmelCase		:     List[Any]=None      , ):
      _UpperCAmelCase         =					parent
      _UpperCAmelCase         =					batch_size
      _UpperCAmelCase         =					seq_length
      _UpperCAmelCase         =					is_training
      _UpperCAmelCase         =					use_input_lengths
      _UpperCAmelCase         =					use_token_type_ids
      _UpperCAmelCase         =					use_labels
      _UpperCAmelCase         =					gelu_activation
      _UpperCAmelCase         =					sinusoidal_embeddings
      _UpperCAmelCase         =					causal
      _UpperCAmelCase         =					asm
      _UpperCAmelCase         =					n_langs
      _UpperCAmelCase         =					vocab_size
      _UpperCAmelCase         =					n_special
      _UpperCAmelCase         =					hidden_size
      _UpperCAmelCase         =					num_hidden_layers
      _UpperCAmelCase         =					num_attention_heads
      _UpperCAmelCase         =					hidden_dropout_prob
      _UpperCAmelCase         =					attention_probs_dropout_prob
      _UpperCAmelCase         =					max_position_embeddings
      _UpperCAmelCase         =					type_vocab_size
      _UpperCAmelCase         =					type_sequence_label_size
      _UpperCAmelCase         =					initializer_range
      _UpperCAmelCase         =					num_labels
      _UpperCAmelCase         =					num_choices
      _UpperCAmelCase         =					summary_type
      _UpperCAmelCase         =					use_proj
      _UpperCAmelCase         =					scope


    def    lowerCAmelCase_     (    self		:     Tuple    ):
      _UpperCAmelCase         =					ids_tensor([self.batch_size, self.seq_length]      , self.vocab_size    )
      _UpperCAmelCase         =					random_attention_mask([self.batch_size, self.seq_length]    )

      _UpperCAmelCase         =					None
      if self.use_input_lengths:
        _UpperCAmelCase         =					(
            ids_tensor([self.batch_size]      , vocab_size=2    ) + self.seq_length - 2
        )  # small variation of seq_length

      _UpperCAmelCase         =					None
      if self.use_token_type_ids:
        _UpperCAmelCase         =					ids_tensor([self.batch_size, self.seq_length]      , self.n_langs    )

      _UpperCAmelCase         =					None
      _UpperCAmelCase         =					None
      _UpperCAmelCase         =					None
      if self.use_labels:
        _UpperCAmelCase         =					ids_tensor([self.batch_size]      , self.type_sequence_label_size    )
        _UpperCAmelCase         =					ids_tensor([self.batch_size, self.seq_length]      , self.num_labels    )
        _UpperCAmelCase         =					ids_tensor([self.batch_size]      , 2    ).float()
        _UpperCAmelCase         =					ids_tensor([self.batch_size]      , self.num_choices    )

      _UpperCAmelCase         =					self.get_config()

      return (
          config,
          input_ids,
          token_type_ids,
          input_lengths,
          sequence_labels,
          token_labels,
          is_impossible_labels,
          choice_labels,
          input_mask,
      )


    def    lowerCAmelCase_     (    self		:     Tuple    ):
      return FlaubertConfig(
          vocab_size=self.vocab_size      , n_special=self.n_special      , emb_dim=self.hidden_size      , n_layers=self.num_hidden_layers      , n_heads=self.num_attention_heads      , dropout=self.hidden_dropout_prob      , attention_dropout=self.attention_probs_dropout_prob      , gelu_activation=self.gelu_activation      , sinusoidal_embeddings=self.sinusoidal_embeddings      , asm=self.asm      , causal=self.causal      , n_langs=self.n_langs      , max_position_embeddings=self.max_position_embeddings      , initializer_range=self.initializer_range      , summary_type=self.summary_type      , use_proj=self.use_proj      , )


    def    lowerCAmelCase_     (    self		:     Optional[Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     Tuple      , ):
      _UpperCAmelCase         =					FlaubertModel(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(__lowerCAmelCase      , lengths=__lowerCAmelCase      , langs=__lowerCAmelCase    )
      _UpperCAmelCase         =					model(__lowerCAmelCase      , langs=__lowerCAmelCase    )
      _UpperCAmelCase         =					model(__lowerCAmelCase    )
      self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)    )


    def    lowerCAmelCase_     (    self		:     int      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     List[str]      , ):
      _UpperCAmelCase         =					FlaubertWithLMHeadModel(__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()

      _UpperCAmelCase         =					model(__lowerCAmelCase      , token_type_ids=__lowerCAmelCase      , labels=__lowerCAmelCase    )
      self.parent.assertEqual(result.loss.shape      , ()    )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.vocab_size)    )


    def    lowerCAmelCase_     (    self		:     List[Any]      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     int      , ):
      _UpperCAmelCase         =					FlaubertForQuestionAnsweringSimple(__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()

      _UpperCAmelCase         =					model(__lowerCAmelCase    )

      _UpperCAmelCase         =					model(__lowerCAmelCase      , start_positions=__lowerCAmelCase      , end_positions=__lowerCAmelCase    )
      self.parent.assertEqual(result.start_logits.shape      , (self.batch_size, self.seq_length)    )
      self.parent.assertEqual(result.end_logits.shape      , (self.batch_size, self.seq_length)    )


    def    lowerCAmelCase_     (    self		:     Optional[Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     List[str]      , ):
      _UpperCAmelCase         =					FlaubertForQuestionAnswering(__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()

      _UpperCAmelCase         =					model(__lowerCAmelCase    )

      _UpperCAmelCase         =					model(
          __lowerCAmelCase      , start_positions=__lowerCAmelCase      , end_positions=__lowerCAmelCase      , cls_index=__lowerCAmelCase      , is_impossible=__lowerCAmelCase      , p_mask=__lowerCAmelCase      , )

      _UpperCAmelCase         =					model(
          __lowerCAmelCase      , start_positions=__lowerCAmelCase      , end_positions=__lowerCAmelCase      , cls_index=__lowerCAmelCase      , is_impossible=__lowerCAmelCase      , )

      ((_UpperCAmelCase)       , )         =					result_with_labels.to_tuple()

      _UpperCAmelCase         =					model(__lowerCAmelCase      , start_positions=__lowerCAmelCase      , end_positions=__lowerCAmelCase    )

      ((_UpperCAmelCase)       , )         =					result_with_labels.to_tuple()

      self.parent.assertEqual(result_with_labels.loss.shape      , ()    )
      self.parent.assertEqual(result.start_top_log_probs.shape      , (self.batch_size, model.config.start_n_top)    )
      self.parent.assertEqual(result.start_top_index.shape      , (self.batch_size, model.config.start_n_top)    )
      self.parent.assertEqual(
          result.end_top_log_probs.shape      , (self.batch_size, model.config.start_n_top * model.config.end_n_top)    )
      self.parent.assertEqual(
          result.end_top_index.shape      , (self.batch_size, model.config.start_n_top * model.config.end_n_top)    )
      self.parent.assertEqual(result.cls_logits.shape      , (self.batch_size,)    )


    def    lowerCAmelCase_     (    self		:     List[Any]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     Tuple      , ):
      _UpperCAmelCase         =					FlaubertForSequenceClassification(__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()

      _UpperCAmelCase         =					model(__lowerCAmelCase    )
      _UpperCAmelCase         =					model(__lowerCAmelCase      , labels=__lowerCAmelCase    )

      self.parent.assertEqual(result.loss.shape      , ()    )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.type_sequence_label_size)    )


    def    lowerCAmelCase_     (    self		:     List[str]      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Optional[Any]      , ):
      _UpperCAmelCase         =					self.num_labels
      _UpperCAmelCase         =					FlaubertForTokenClassification(__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()

      _UpperCAmelCase         =					model(__lowerCAmelCase      , attention_mask=__lowerCAmelCase      , labels=__lowerCAmelCase    )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.num_labels)    )


    def    lowerCAmelCase_     (    self		:     List[Any]      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     Tuple      , ):
      _UpperCAmelCase         =					self.num_choices
      _UpperCAmelCase         =					FlaubertForMultipleChoice(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					input_ids.unsqueeze(1    ).expand(-1      , self.num_choices      , -1    ).contiguous()
      _UpperCAmelCase         =					token_type_ids.unsqueeze(1    ).expand(-1      , self.num_choices      , -1    ).contiguous()
      _UpperCAmelCase         =					input_mask.unsqueeze(1    ).expand(-1      , self.num_choices      , -1    ).contiguous()
      _UpperCAmelCase         =					model(
          __lowerCAmelCase      , attention_mask=__lowerCAmelCase      , token_type_ids=__lowerCAmelCase      , labels=__lowerCAmelCase      , )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.num_choices)    )


    def    lowerCAmelCase_     (    self		:     int    ):
      _UpperCAmelCase         =					self.prepare_config_and_inputs()
      (
          (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , (
          _UpperCAmelCase
      )       , 
      )         =					config_and_inputs
      _UpperCAmelCase         =					{
          """input_ids""": input_ids,
          """token_type_ids""": token_type_ids,
          """lengths""": input_lengths,
          """attention_mask""": input_mask,
      }
      return config, inputs_dict



@require_torch
class 				a (						lowerCAmelCase_		,							lowerCAmelCase_		,							unittest.TestCase      ):
    _snake_case   :     List[str]   							=							(
        (
            FlaubertModel,
            FlaubertWithLMHeadModel,
            FlaubertForQuestionAnswering,
            FlaubertForQuestionAnsweringSimple,
            FlaubertForSequenceClassification,
            FlaubertForTokenClassification,
            FlaubertForMultipleChoice,
        )
        if is_torch_available()
        else ()
    )
    _snake_case   :     str   							=							(
        {
            'feature-extraction': FlaubertModel,
            'fill-mask': FlaubertWithLMHeadModel,
            'question-answering': FlaubertForQuestionAnsweringSimple,
            'text-classification': FlaubertForSequenceClassification,
            'token-classification': FlaubertForTokenClassification,
            'zero-shot': FlaubertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )


    def    lowerCAmelCase_     (    self		:     Dict      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     Optional[Any]      , __lowerCAmelCase		:     List[Any]    ):
      if (
          pipeline_test_casse_name == "QAPipelineTests"
          and tokenizer_name is not None
          and not tokenizer_name.endswith("""Fast"""    )
      ):
        # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
        # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
        # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
        return True

      return False


    def    lowerCAmelCase_     (    self		:     Any      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     List[Any]=False    ):
      _UpperCAmelCase         =					super()._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase      , return_labels=__lowerCAmelCase    )

      if return_labels:
        if model_class.__name__ == "FlaubertForQuestionAnswering":
          _UpperCAmelCase         =					torch.zeros(
              self.model_tester.batch_size      , dtype=torch.long      , device=__lowerCAmelCase    )
          _UpperCAmelCase         =					torch.zeros(
              self.model_tester.batch_size      , dtype=torch.long      , device=__lowerCAmelCase    )

      return inputs_dict


    def    lowerCAmelCase_     (    self		:     str    ):
      _UpperCAmelCase         =					FlaubertModelTester(self    )
      _UpperCAmelCase         =					ConfigTester(self      , config_class=__lowerCAmelCase      , emb_dim=37    )


    def    lowerCAmelCase_     (    self		:     int    ):
      self.config_tester.run_common_tests()


    def    lowerCAmelCase_     (    self		:     Tuple    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_model(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     str    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_lm_head(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Dict    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_simple_qa(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_qa(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_sequence_classif(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_token_classif(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     int    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_flaubert_multiple_choice(*__lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					FlaubertModel.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )


    @slow
    @require_torch_gpu
    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase       , _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
        # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
        if model_class == FlaubertForMultipleChoice:
          return

        _UpperCAmelCase         =					True
        _UpperCAmelCase         =					model_class(config=__lowerCAmelCase    )

        _UpperCAmelCase         =					self._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase    )
        _UpperCAmelCase         =					torch.jit.trace(
            __lowerCAmelCase      , (inputs_dict["""input_ids"""].to("""cpu"""    ), inputs_dict["""attention_mask"""].to("""cpu"""    ))    )

        with tempfile.TemporaryDirectory() as tmp:
          torch.jit.save(__lowerCAmelCase      , os.path.join(__lowerCAmelCase      , """traced_model.pt"""    )    )
          _UpperCAmelCase         =					torch.jit.load(os.path.join(__lowerCAmelCase      , """traced_model.pt"""    )      , map_location=__lowerCAmelCase    )
          loaded(inputs_dict["""input_ids"""].to(__lowerCAmelCase    )      , inputs_dict["""attention_mask"""].to(__lowerCAmelCase    )    )



@require_torch
class 				a (						unittest.TestCase      ):


    @slow
    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      _UpperCAmelCase         =					FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased"""    )
      _UpperCAmelCase         =					torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
      with torch.no_grad():
        _UpperCAmelCase         =					model(__lowerCAmelCase    )[0]
      _UpperCAmelCase         =					torch.Size((1, 11, 768)    )
      self.assertEqual(output.shape      , __lowerCAmelCase    )
      _UpperCAmelCase         =					torch.tensor(
          [[[-2.6_251, -1.4_298, -0.0_227], [-2.8_510, -1.6_387, 0.2_258], [-2.8_114, -1.1_832, -0.3_066]]]    )

      self.assertTrue(torch.allclose(output[:, :3, :3]      , __lowerCAmelCase      , atol=1e-4    )    )

"""simple docstring"""


from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
    from .deepspeed import (
        DeepSpeedEngineWrapper,
        DeepSpeedOptimizerWrapper,
        DeepSpeedSchedulerWrapper,
        DummyOptim,
        DummyScheduler,
        HfDeepSpeedConfig,
    )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers
