


"""simple docstring"""




# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
	import jax
	import jaxlib

__UpperCamelCase						=      get_logger()

__UpperCamelCase						=      None
class 					lowerCAmelCase	(					TensorFormatter[Mapping, """jax.Array""", Mapping]    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__=None						,       lowerCAmelCase__=None						,       **lowerCAmelCase__				)     ->							str:
								super().__init__(features=lowerCAmelCase__				)
								import jax
								from jaxlib.xla_client import Device

								if isinstance(lowerCAmelCase__						,       lowerCAmelCase__				):
														raise ValueError(
														    F'Expected {device} to be a `str` not {type(lowerCAmelCase__				)}, as `jaxlib.xla_extension.Device` '
														    'is not serializable neither with `pickle` nor with `dill`. Instead you can surround '
														    'the device with `str()` to get its string identifier that will be internally mapped '
														    'to the actual `jaxlib.xla_extension.Device`.'				)
								SCREAMING_SNAKE_CASE													=					device if isinstance(lowerCAmelCase__						,       lowerCAmelCase__				) else str(jax.devices()[0]				)
								# using global variable since `jaxlib.xla_extension.Device` is not serializable neither
								# with `pickle` nor with `dill`, so we need to use a global variable instead
								global DEVICE_MAPPING
								if DEVICE_MAPPING is None:
														SCREAMING_SNAKE_CASE													=					self._map_devices_to_str()
								if self.device not in list(DEVICE_MAPPING.keys()				):
														logger.warning(
														    F'Device with string identifier {self.device} not listed among the available '
														    F'devices: {list(DEVICE_MAPPING.keys()				)}, so falling back to the default '
														    F'device: {str(jax.devices()[0]				)}.'				)
														SCREAMING_SNAKE_CASE													=					str(jax.devices()[0]				)
								SCREAMING_SNAKE_CASE													=					jnp_array_kwargs







		@staticmethod
		def 							__A						(	)     ->							Dict[str, "jaxlib.xla_extension.Device"]:
								import jax

								return {str(lowerCAmelCase__				): device for device in jax.devices()}







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							List[Any]:
								import jax
								import jax.numpy as jnp

								if isinstance(lowerCAmelCase__						,       lowerCAmelCase__				) and column:
														if all(
														    isinstance(lowerCAmelCase__						,       jax.Array				) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column				):
																				return jnp.stack(lowerCAmelCase__						,       axis=0				)
								return column







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Optional[Any]:
								import jax
								import jax.numpy as jnp

								if isinstance(lowerCAmelCase__						,       (str, bytes, type(lowerCAmelCase__				))				):
														return value
								elif isinstance(lowerCAmelCase__						,       (np.character, np.ndarray)				) and np.issubdtype(value.dtype						,       np.character				):
														return value.tolist()

								SCREAMING_SNAKE_CASE													=					{}

								if isinstance(lowerCAmelCase__						,       (np.number, np.ndarray)				) and np.issubdtype(value.dtype						,       np.integer				):
														# the default int precision depends on the jax config
														# see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
														if jax.config.jax_enable_xaa:
																				SCREAMING_SNAKE_CASE													=					{'dtype': jnp.intaa}
														else:
																				SCREAMING_SNAKE_CASE													=					{'dtype': jnp.intaa}
								elif isinstance(lowerCAmelCase__						,       (np.number, np.ndarray)				) and np.issubdtype(value.dtype						,       np.floating				):
														SCREAMING_SNAKE_CASE													=					{'dtype': jnp.floataa}
								elif config.PIL_AVAILABLE and "PIL" in sys.modules:
														import PIL.Image

														if isinstance(lowerCAmelCase__						,       PIL.Image.Image				):
																				SCREAMING_SNAKE_CASE													=					np.asarray(lowerCAmelCase__				)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
								global DEVICE_MAPPING
								if DEVICE_MAPPING is None:
														SCREAMING_SNAKE_CASE													=					self._map_devices_to_str()

								with jax.default_device(DEVICE_MAPPING[self.device]				):
														# calling jnp.array on a np.ndarray does copy the data
														# see https://github.com/google/jax/issues/4486
														return jnp.array(lowerCAmelCase__						,       **{**default_dtype, **self.jnp_array_kwargs}				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							str:
								import jax

								# support for torch, tf, jax etc.
								if config.TORCH_AVAILABLE and "torch" in sys.modules:
														import torch

														if isinstance(lowerCAmelCase__						,       torch.Tensor				):
																				return self._tensorize(data_struct.detach().cpu().numpy()[()]				)
								if hasattr(lowerCAmelCase__						,       '__array__'				) and not isinstance(lowerCAmelCase__						,       jax.Array				):
														SCREAMING_SNAKE_CASE													=					data_struct.__array__()
								# support for nested types like struct of list of struct
								if isinstance(lowerCAmelCase__						,       np.ndarray				):
														if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
																				return self._consolidate([self.recursive_tensorize(lowerCAmelCase__				) for substruct in data_struct]				)
								elif isinstance(lowerCAmelCase__						,       (list, tuple)				):
														return self._consolidate([self.recursive_tensorize(lowerCAmelCase__				) for substruct in data_struct]				)
								return self._tensorize(lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Union[str, Any]:
								return map_nested(self._recursive_tensorize						,       lowerCAmelCase__						,       map_list=lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Mapping:
								SCREAMING_SNAKE_CASE													=					self.numpy_arrow_extractor().extract_row(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.python_features_decoder.decode_row(lowerCAmelCase__				)
								return self.recursive_tensorize(lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							"jax.Array":
								SCREAMING_SNAKE_CASE													=					self.numpy_arrow_extractor().extract_column(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.python_features_decoder.decode_column(lowerCAmelCase__						,       pa_table.column_names[0]				)
								SCREAMING_SNAKE_CASE													=					self.recursive_tensorize(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self._consolidate(lowerCAmelCase__				)
								return column







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Mapping:
								SCREAMING_SNAKE_CASE													=					self.numpy_arrow_extractor().extract_batch(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.python_features_decoder.decode_batch(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.recursive_tensorize(lowerCAmelCase__				)
								for column_name in batch:
														SCREAMING_SNAKE_CASE													=					self._consolidate(batch[column_name]				)
								return batch






"""simple docstring"""




import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

__UpperCamelCase						=      [
    '''cross_validation.py''',
    '''gradient_accumulation.py''',
    '''local_sgd.py''',
    '''multi_process_metrics.py''',
    '''memory.py''',
    '''automatic_gradient_accumulation.py''',
    '''fsdp_with_peak_mem_tracking.py''',
    '''deepspeed_with_config_support.py''',
    '''megatron_lm_gpt_pretraining.py''',
]
class 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None				)     ->							int:
								SCREAMING_SNAKE_CASE													=					None
								SCREAMING_SNAKE_CASE													=					os.path.abspath(os.path.join('examples'						,       'by_feature'				)				)
								SCREAMING_SNAKE_CASE													=					os.path.abspath('examples'				)
								for item in os.listdir(lowerCAmelCase__				):
														if item not in EXCLUDE_EXAMPLES:
																				SCREAMING_SNAKE_CASE													=					os.path.join(lowerCAmelCase__						,       lowerCAmelCase__				)
																				if os.path.isfile(lowerCAmelCase__				) and ".py" in item_path:
																										with self.subTest(
																										    tested_script=lowerCAmelCase__						,       feature_script=lowerCAmelCase__						,       tested_section='main()' if parser_only else 'training_function()'						,       ):
																																SCREAMING_SNAKE_CASE													=					compare_against_test(
																																    os.path.join(lowerCAmelCase__						,       lowerCAmelCase__				)						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)
																																SCREAMING_SNAKE_CASE													=					'\n'.join(lowerCAmelCase__				)
																																if special_strings is not None:
																																						for string in special_strings:
																																												SCREAMING_SNAKE_CASE													=					diff.replace(lowerCAmelCase__						,       ''				)
																																self.assertEqual(lowerCAmelCase__						,       ''				)







		def 							__A						(	self				)     ->							Optional[int]:
								self.one_complete_example('complete_nlp_example.py'						,       lowerCAmelCase__				)
								self.one_complete_example('complete_nlp_example.py'						,       lowerCAmelCase__				)







		def 							__A						(	self				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					os.path.abspath(os.path.join('examples'						,       'cv_example.py'				)				)
								SCREAMING_SNAKE_CASE													=					[
								    ' ' * 16 + '{\n\n',
								    ' ' * 20 + '"accuracy": eval_metric["accuracy"],\n\n',
								    ' ' * 20 + '"f1": eval_metric["f1"],\n\n',
								    ' ' * 20 + '"train_loss": total_loss.item() / len(train_dataloader),\n\n',
								    ' ' * 20 + '"epoch": epoch,\n\n',
								    ' ' * 16 + '},\n\n',
								    ' ' * 16 + 'step=epoch,\n',
								    ' ' * 12,
								    ' ' * 8 + 'for step, batch in enumerate(active_dataloader):\n',
								]
								self.one_complete_example('complete_cv_example.py'						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)
								self.one_complete_example('complete_cv_example.py'						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)





@mock.patch.dict(os.environ	,    {"""TESTING_MOCKED_DATALOADERS""": """1"""}    )
class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     List[str]								=     False







		@classmethod
		def 							__A						(	cls				)     ->							List[str]:
								super().setUpClass()
								SCREAMING_SNAKE_CASE													=					tempfile.mkdtemp()
								SCREAMING_SNAKE_CASE													=					os.path.join(cls._tmpdir						,       'default_config.yml'				)

								write_basic_config(save_location=cls.configPath				)
								SCREAMING_SNAKE_CASE													=					['accelerate', 'launch', '--config_file', cls.configPath]







		@classmethod
		def 							__A						(	cls				)     ->							Dict:
								super().tearDownClass()
								shutil.rmtree(cls._tmpdir				)







		def 							__A						(	self				)     ->							int:
								SCREAMING_SNAKE_CASE													=					F'\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps epoch\n        --output_dir {self.tmpdir}\n        '.split()
								run_command(self._launch_args + testargs				)
								self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,       'epoch_0'				)				)				)







		def 							__A						(	self				)     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					F'\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        '.split()
								SCREAMING_SNAKE_CASE													=					run_command(self._launch_args + testargs				)
								self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,       'step_2'				)				)				)







		def 							__A						(	self				)     ->							Union[str, Any]:
								SCREAMING_SNAKE_CASE													=					F'\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir						,       "epoch_0"				)}\n        '.split()
								SCREAMING_SNAKE_CASE													=					run_command(self._launch_args + testargs						,       return_stdout=lowerCAmelCase__				)
								self.assertNotIn('epoch 0:'						,       lowerCAmelCase__				)
								self.assertIn('epoch 1:'						,       lowerCAmelCase__				)







		def 							__A						(	self				)     ->							int:
								SCREAMING_SNAKE_CASE													=					F'\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir						,       "step_2"				)}\n        '.split()
								SCREAMING_SNAKE_CASE													=					run_command(self._launch_args + testargs						,       return_stdout=lowerCAmelCase__				)
								if torch.cuda.is_available():
														SCREAMING_SNAKE_CASE													=					torch.cuda.device_count()
								else:
														SCREAMING_SNAKE_CASE													=					1
								if num_processes > 1:
														self.assertNotIn('epoch 0:'						,       lowerCAmelCase__				)
														self.assertIn('epoch 1:'						,       lowerCAmelCase__				)
								else:
														self.assertIn('epoch 0:'						,       lowerCAmelCase__				)
														self.assertIn('epoch 1:'						,       lowerCAmelCase__				)







		@slow
		def 							__A						(	self				)     ->							Optional[Any]:
								SCREAMING_SNAKE_CASE													=					'\n        examples/by_feature/cross_validation.py\n        --num_folds 2\n        '.split()
								with mock.patch.dict(os.environ						,       {'TESTING_MOCKED_DATALOADERS': '0'}				):
														SCREAMING_SNAKE_CASE													=					run_command(self._launch_args + testargs						,       return_stdout=lowerCAmelCase__				)
														SCREAMING_SNAKE_CASE													=					re.findall('({.+})'						,       lowerCAmelCase__				)
														SCREAMING_SNAKE_CASE													=					[r for r in results if 'accuracy' in r][-1]
														SCREAMING_SNAKE_CASE													=					ast.literal_eval(lowerCAmelCase__				)
														self.assertGreaterEqual(results['accuracy']						,       0.75				)







		def 							__A						(	self				)     ->							str:
								SCREAMING_SNAKE_CASE													=					['examples/by_feature/multi_process_metrics.py']
								run_command(self._launch_args + testargs				)







		@require_trackers
		@mock.patch.dict(os.environ						,       {'WANDB_MODE': 'offline'}				)
		def 							__A						(	self				)     ->							Union[str, Any]:
								with tempfile.TemporaryDirectory() as tmpdir:
														SCREAMING_SNAKE_CASE													=					F'\n            examples/by_feature/tracking.py\n            --with_tracking\n            --project_dir {tmpdir}\n            '.split()
														run_command(self._launch_args + testargs				)
														self.assertTrue(os.path.exists(os.path.join(lowerCAmelCase__						,       'tracking'				)				)				)







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					['examples/by_feature/gradient_accumulation.py']
								run_command(self._launch_args + testargs				)







		def 							__A						(	self				)     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					['examples/by_feature/local_sgd.py']
								run_command(self._launch_args + testargs				)


