




'''simple docstring'''


import bisect
import copy
from collections import defaultdict

import numpy as np
from torch.utils.data import BatchSampler, Sampler

from utils import logger
def 					SCREAMING_SNAKE_CASE(   __lowercase		,				__lowercase	)			-> List[Any]:
     A:		List[Any]          =		copy.deepcopy(__lowercase	)
     A:		str          =		sorted(__lowercase	)
     A:		Tuple          =		[bisect.bisect_right(__lowercase		,				__lowercase	) for y in x]
     return quantized
def 					SCREAMING_SNAKE_CASE(   __lowercase		,				__lowercase=0	)			-> Union[str, Any]:
     A:		Dict          =		np.arange(start=3		,				stop=__lowercase		,				step=4	).tolist() if k > 0 else [1_0]
     A:		Optional[int]          =		_quantize(__lowercase		,				__lowercase	)
     # count number of elements per group
     A:		Dict          =		np.unique(__lowercase		,				return_counts=__lowercase	)[1]
     A:		Tuple          =		[0] + bins + [np.inf]
     logger.info('''Using {} as bins for aspect lengths quantization'''.format(__lowercase	)	)
     logger.info('''Count of instances per bin: {}'''.format(__lowercase	)	)
     return groups






class 						lowerCAmelCase_	(       UpperCAmelCase_							):







  '''simple docstring'''

  def __init__( self				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							Optional[int]	)			->		List[str]:







       '''simple docstring'''



       if not isinstance(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	):
            raise ValueError(
                '''sampler should be an instance of torch.utils.data.Sampler, but got sampler={}'''.format(SCREAMING_SNAKE_CASE_	)	)
       A:		Union[str, Any]          =		sampler
       A:		str          =		group_ids
       A:		List[str]          =		batch_size



  def __iter__( self				:							List[Any]	)			->		Dict:







       '''simple docstring'''



       A:		Optional[Any]          =		defaultdict(SCREAMING_SNAKE_CASE_	)
       A:		Tuple          =		defaultdict(SCREAMING_SNAKE_CASE_	)

       A:		Optional[int]          =		0
       for idx in self.sampler:
            A:		List[Any]          =		self.group_ids[idx]
            buffer_per_group[group_id].append(SCREAMING_SNAKE_CASE_	)
            samples_per_group[group_id].append(SCREAMING_SNAKE_CASE_	)
            if len(buffer_per_group[group_id]	) == self.batch_size:
                 yield buffer_per_group[group_id]  # TODO
                 num_batches += 1
                 del buffer_per_group[group_id]
            assert len(buffer_per_group[group_id]	) < self.batch_size

       # now we have run out of elements that satisfy
       # the group criteria, let's return the remaining
       # elements so that the size of the sampler is
       # deterministic
       A:		Tuple          =		len(self	)
       A:		Any          =		expected_num_batches - num_batches
       if num_remaining > 0:
            # for the remaining batches, group the batches by similar lengths
            A:		Any          =		[]
            for group_id, idxs in sorted(buffer_per_group.items()					,  key=lambda SCREAMING_SNAKE_CASE_	: x[0]	):
                 batch_idx.extend(SCREAMING_SNAKE_CASE_	)
                 if len(SCREAMING_SNAKE_CASE_	) >= self.batch_size:
                      yield batch_idx[: self.batch_size]
                      A:		List[Any]          =		batch_idx[self.batch_size :]
                      num_remaining -= 1
            if len(SCREAMING_SNAKE_CASE_	) > 0:
                 yield batch_idx
                 num_remaining -= 1
       assert num_remaining == 0




  def __len__( self				:							Optional[int]	)			->		Optional[int]:







       '''simple docstring'''



       return (len(self.sampler	) + self.batch_size - 1) // self.batch_size












'''simple docstring'''


import importlib
import importlib.metadata
import os
import warnings
from distutils.util import strtobool
from functools import lru_cache

import torch
from packaging import version
from packaging.version import parse

from .environment import parse_flag_from_env
from .versions import compare_versions, is_torch_version


try:
     import torch_xla.core.xla_model as xm  # noqa: F401

     UpperCamelCase			     =					True
except ImportError:
     UpperCamelCase			     =					False


# Cache this result has it's a C FFI call which can be pretty time-consuming
UpperCamelCase			     =					torch.distributed.is_available()
def 					SCREAMING_SNAKE_CASE(   __lowercase	)			-> Optional[int]:
     # Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version
     A:		Optional[Any]          =		importlib.util.find_spec(__lowercase	) is not None
     if package_exists:
          try:
               A:		Dict          =		importlib.metadata.metadata(__lowercase	)
               return True
          except importlib.metadata.PackageNotFoundError:
               return False
def 					SCREAMING_SNAKE_CASE(   )			-> bool:
     return _torch_distributed_available
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[int]:
     try:
          pass
     except ImportError:
          print(
              '''Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not'''
              ''' detected. If you see \"ValueError: Invalid backend: \'ccl\'\" error, please install Intel(R) oneCCL'''
              ''' Bindings for PyTorch*.'''	)
     return (
         importlib.util.find_spec('''torch_ccl'''	) is not None
         or importlib.util.find_spec('''oneccl_bindings_for_pytorch'''	) is not None
     )
def 					SCREAMING_SNAKE_CASE(   )			-> List[str]:
     return importlib.metadata.version('''oneccl_bind_pt'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''transformer_engine'''	)
@lru_cache
def 					SCREAMING_SNAKE_CASE(   __lowercase=True	)			-> Union[str, Any]:
     # Due to bugs on the amp series GPUs, we disable torch-xla on them
     if torch.cuda.is_available():
          return False
     if _tpu_available and check_device:
          try:
               # Will raise a RuntimeError if no XLA configuration is found
               A:		Tuple          =		xm.xla_device()
               return True
          except RuntimeError:
               return False
     return _tpu_available
def 					SCREAMING_SNAKE_CASE(   )			-> List[str]:
     return _is_package_available('''deepspeed'''	)
def 					SCREAMING_SNAKE_CASE(   __lowercase=False	)			-> Any:
     if is_tpu_available():
          return not ignore_tpu
     if torch.cuda.is_available():
          return torch.cuda.is_bfaa_supported()
     if is_npu_available():
          return False
     return True
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     A:		List[str]          =		_is_package_available('''bitsandbytes'''	)
     if package_exists:
          A:		Union[str, Any]          =		version.parse(importlib.metadata.version('''bitsandbytes'''	)	)
          return compare_versions(__lowercase		,				'''>='''		,				'''0.39.0'''	)
     return False
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     A:		int          =		_is_package_available('''bitsandbytes'''	)
     if package_exists:
          A:		Tuple          =		version.parse(importlib.metadata.version('''bitsandbytes'''	)	)
          return compare_versions(__lowercase		,				'''>='''		,				'''0.37.2'''	)
     return False
def 					SCREAMING_SNAKE_CASE(   )			-> int:
     return _is_package_available('''bitsandbytes'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[int]:
     if strtobool(os.environ.get('''ACCELERATE_USE_MEGATRON_LM'''		,				'''False'''	)	) == 1:
          A:		Any          =		importlib.util.find_spec('''megatron'''	) is not None
          if package_exists:
               try:
                    A:		str          =		parse(importlib.metadata.version('''megatron-lm'''	)	)
                    return compare_versions(__lowercase		,				'''>='''		,				'''2.2.0'''	)
               except Exception as e:
                    warnings.warn(F"""Parse Megatron version failed. Exception:{e}"""	)
                    return False
def 					SCREAMING_SNAKE_CASE(   )			-> Any:
     return _is_package_available('''safetensors'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Tuple:
     return _is_package_available('''transformers'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''datasets'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''aim'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> List[Any]:
     return _is_package_available('''tensorboard'''	) or _is_package_available('''tensorboardX'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Dict:
     return _is_package_available('''wandb'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[Any]:
     return _is_package_available('''comet_ml'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Any:
     return _is_package_available('''boto3'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     if _is_package_available('''rich'''	):
          if "ACCELERATE_DISABLE_RICH" in os.environ:
               warnings.warn(
                   '''`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`.'''	)
               return not parse_flag_from_env('''ACCELERATE_DISABLE_RICH'''		,				__lowercase	)
          return parse_flag_from_env('''ACCELERATE_ENABLE_RICH'''		,				__lowercase	)
     return False
def 					SCREAMING_SNAKE_CASE(   )			-> int:
     return _is_package_available('''sagemaker'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''tqdm'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     return _is_package_available('''mlflow'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> int:
     return is_torch_version('''>='''		,				'''1.12'''	) and torch.backends.mps.is_available() and torch.backends.mps.is_built()
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     def get_major_and_minor_from_version(__lowercase	):
          return str(version.parse(__lowercase	).major	) + "." + str(version.parse(__lowercase	).minor	)

     A:		int          =		importlib.metadata.version('''torch'''	)
     if importlib.util.find_spec('''intel_extension_for_pytorch'''	) is None:
          return False
     A:		Optional[Any]          =		'''N/A'''
     try:
          A:		List[Any]          =		importlib.metadata.version('''intel_extension_for_pytorch'''	)
     except importlib.metadata.PackageNotFoundError:
          return False
     A:		Optional[int]          =		get_major_and_minor_from_version(_torch_version	)
     A:		List[str]          =		get_major_and_minor_from_version(_ipex_version	)
     if torch_major_and_minor != ipex_major_and_minor:
          warnings.warn(
              F"""Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,"""
              F""" but PyTorch {_torch_version} is found. Please switch to the matching version and run again."""	)
          return False
     return True
@lru_cache
def 					SCREAMING_SNAKE_CASE(   __lowercase=False	)			-> Optional[Any]:
     if importlib.util.find_spec('''torch'''	) is None or importlib.util.find_spec('''torch_npu'''	) is None:
          return False

     import torch
     import torch_npu  # noqa: F401

     if check_device:
          try:
               # Will raise a RuntimeError if no NPU is found
               A:		Tuple          =		torch.npu.device_count()
               return torch.npu.is_available()
          except RuntimeError:
               return False
     return hasattr(__lowercase		,				'''npu'''	) and torch.npu.is_available()


@lru_cache
def 					SCREAMING_SNAKE_CASE(   __lowercase=False	)			-> Tuple:
     if not parse_flag_from_env('''ACCELERATE_USE_XPU'''		,				default=__lowercase	):
          return False
     if is_ipex_available():
          import torch

          if is_torch_version('''<='''		,				'''1.12'''	):
               return False
     else:
          return False

     import intel_extension_for_pytorch  # noqa: F401

     if check_device:
          try:
               # Will raise a RuntimeError if no XPU  is found
               A:		str          =		torch.xpu.device_count()
               return torch.xpu.is_available()
          except RuntimeError:
               return False
     return hasattr(__lowercase		,				'''xpu'''	) and torch.xpu.is_available()






