

from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor





class A__    :



 """simple docstring"""







 def __init__(		self :	Dict     ,						A_ :	int     ,						A_ :	List[str]=1_3     ,						A_ :	List[Any]=3_0     ,						A_ :	Tuple=2     ,						A_ :	Optional[Any]=3     ,						A_ :	Union[str, Any]=True     ,						A_ :	Optional[Any]=True     ,						A_ :	Any=3_2     ,						A_ :	Any=2     ,						A_ :	int=4     ,						A_ :	Dict=3_7     ,						A_ :	int="gelu"     ,						A_ :	str=0.1     ,						A_ :	int=0.1     ,						A_ :	Any=1_0     ,						A_ :	List[str]=0.02     ,						A_ :	int=3     ,						A_ :	Optional[int]=0.6     ,						A_ :	str=None     ,						):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    parent
      _lowerCAmelCase   :					Dict        =    batch_size
      _lowerCAmelCase   :					Dict        =    image_size
      _lowerCAmelCase   :					int        =    patch_size
      _lowerCAmelCase   :					str        =    num_channels
      _lowerCAmelCase   :					Optional[Any]        =    is_training
      _lowerCAmelCase   :					Union[str, Any]        =    use_labels
      _lowerCAmelCase   :					str        =    hidden_size
      _lowerCAmelCase   :					Optional[Any]        =    num_hidden_layers
      _lowerCAmelCase   :					List[str]        =    num_attention_heads
      _lowerCAmelCase   :					Tuple        =    intermediate_size
      _lowerCAmelCase   :					Dict        =    hidden_act
      _lowerCAmelCase   :					Dict        =    hidden_dropout_prob
      _lowerCAmelCase   :					List[str]        =    attention_probs_dropout_prob
      _lowerCAmelCase   :					int        =    type_sequence_label_size
      _lowerCAmelCase   :					Union[str, Any]        =    initializer_range
      _lowerCAmelCase   :					int        =    mask_ratio
      _lowerCAmelCase   :					Union[str, Any]        =    scope

      # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
      # (we add 1 for the [CLS] token)
      _lowerCAmelCase   :					Optional[Any]        =    (image_size // patch_size) ** 2
      _lowerCAmelCase   :					Optional[Any]        =    int(math.ceil((1 - mask_ratio) * (num_patches + 1)       )       )

 def    __magic_name__       (		self :	Dict       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

      _lowerCAmelCase   :					Optional[Any]        =    None
      if self.use_labels:
           _lowerCAmelCase   :					Any        =    ids_tensor([self.batch_size]     ,						self.type_sequence_label_size       )

      _lowerCAmelCase   :					int        =    self.get_config()

      return config, pixel_values, labels

 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




      return ViTMAEConfig(
          image_size=self.image_size     ,						patch_size=self.patch_size     ,						num_channels=self.num_channels     ,						hidden_size=self.hidden_size     ,						num_hidden_layers=self.num_hidden_layers     ,						num_attention_heads=self.num_attention_heads     ,						intermediate_size=self.intermediate_size     ,						decoder_hidden_size=self.hidden_size     ,						decoder_num_hidden_layers=self.num_hidden_layers     ,						decoder_num_attention_heads=self.num_attention_heads     ,						decoder_intermediate_size=self.intermediate_size     ,						hidden_act=self.hidden_act     ,						hidden_dropout_prob=self.hidden_dropout_prob     ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,						is_decoder=A_     ,						initializer_range=self.initializer_range     ,						mask_ratio=self.mask_ratio     ,						)

 def    __magic_name__       (		self :	Tuple     ,						A_ :	List[Any]     ,						A_ :	Any     ,						A_ :	Optional[int]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					int        =    TFViTMAEModel(config=A_       )
      _lowerCAmelCase   :					List[str]        =    model(A_     ,						training=A_       )
      self.parent.assertEqual(result.last_hidden_state.shape     ,						(self.batch_size, self.seq_length, self.hidden_size)       )

 def    __magic_name__       (		self :	Tuple     ,						A_ :	Optional[Any]     ,						A_ :	int     ,						A_ :	Any       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    TFViTMAEForPreTraining(A_       )
      _lowerCAmelCase   :					Optional[int]        =    model(A_     ,						training=A_       )
      # expected sequence length = num_patches
      _lowerCAmelCase   :					List[Any]        =    (self.image_size // self.patch_size) ** 2
      _lowerCAmelCase   :					Union[str, Any]        =    self.patch_size**2 * self.num_channels
      self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, num_patches, expected_num_channels)       )

      # test greyscale images
      _lowerCAmelCase   :					Optional[Any]        =    1
      _lowerCAmelCase   :					str        =    TFViTMAEForPreTraining(A_       )

      _lowerCAmelCase   :					Optional[int]        =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
      _lowerCAmelCase   :					Dict        =    model(A_     ,						training=A_       )
      _lowerCAmelCase   :					List[str]        =    self.patch_size**2
      self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, num_patches, expected_num_channels)       )


 def    __magic_name__       (		self :	Union[str, Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    self.prepare_config_and_inputs()
      ((_lowerCAmelCase)      ,			(_lowerCAmelCase)      ,			(_lowerCAmelCase))   :					Dict        =    config_and_inputs
      _lowerCAmelCase   :					Any        =    {"pixel_values": pixel_values}
      return config, inputs_dict





@require_tf
class A__    ( A  ,		A  ,		unittest.TestCase ):



 """simple docstring"""







 _lowercase     :					Tuple      =						(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
 _lowercase     :					List[str]      =						{'''feature-extraction''': TFViTMAEModel} if is_tf_available() else {}

 _lowercase     :					str      =						False
 _lowercase     :					Dict      =						False
 _lowercase     :					Any      =						False
 _lowercase     :					Dict      =						False

 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Union[str, Any]        =    TFViTMAEModelTester(self       )
      _lowerCAmelCase   :					List[str]        =    ConfigTester(self     ,						config_class=A_     ,						has_text_modality=A_     ,						hidden_size=3_7       )

 def    __magic_name__       (		self :	Dict       ):






      '''simple docstring'''




      self.config_tester.run_common_tests()

 @unittest.skip(reason="ViTMAE does not use inputs_embeds"       )
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      pass

 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase      ,			_lowerCAmelCase   :					Union[str, Any]        =    self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _lowerCAmelCase   :					Optional[int]        =    model_class(A_       )
           self.assertIsInstance(model.get_input_embeddings()     ,						(tf.keras.layers.Layer)       )
           _lowerCAmelCase   :					str        =    model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(A_     ,						tf.keras.layers.Layer       )       )

 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase      ,			_lowerCAmelCase   :					str        =    self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _lowerCAmelCase   :					int        =    model_class(A_       )
           _lowerCAmelCase   :					List[str]        =    inspect.signature(model.call       )
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           _lowerCAmelCase   :					Optional[Any]        =    [*signature.parameters.keys()]

           _lowerCAmelCase   :					Optional[Any]        =    ["pixel_values"]
           self.assertListEqual(arg_names[:1]     ,						A_       )

 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A_       )

 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					str        =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*A_       )

 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase      ,			_lowerCAmelCase   :					Optional[int]        =    self.model_tester.prepare_config_and_inputs_for_common()

      _lowerCAmelCase   :					List[Any]        =    int((config.image_size // config.patch_size) ** 2       )
      _lowerCAmelCase   :					Optional[int]        =    np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

      for model_class in self.all_model_classes:
           _lowerCAmelCase   :					List[str]        =    model_class(A_       )
           _lowerCAmelCase   :					str        =    self._prepare_for_class(A_     ,						A_       )

           _lowerCAmelCase   :					int        =    model(A_     ,						noise=A_       )

           _lowerCAmelCase   :					Dict        =    copy.deepcopy(self._prepare_for_class(A_     ,						A_       )       )
           _lowerCAmelCase   :					int        =    model(**A_     ,						noise=A_       )
           _lowerCAmelCase   :					Union[str, Any]        =    outputs_dict[0].numpy()
           _lowerCAmelCase   :					Any        =    outputs_keywords[0].numpy()

           self.assertLess(np.sum(np.abs(output_dict - output_keywords       )       )     ,						1E-6       )

 def    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase      ,			_lowerCAmelCase   :					Union[str, Any]        =    self.model_tester.prepare_config_and_inputs_for_common()

      _lowerCAmelCase   :					List[str]        =    int((config.image_size // config.patch_size) ** 2       )
      _lowerCAmelCase   :					List[str]        =    np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

      def prepare_numpy_arrays(A_ :	str       ):
           _lowerCAmelCase   :					int        =    {}
           for k, v in inputs_dict.items():
                if tf.is_tensor(A_       ):
                     _lowerCAmelCase   :					Tuple        =    v.numpy()
                else:
                     _lowerCAmelCase   :					Optional[Any]        =    np.array(A_       )

           return inputs_np_dict

      for model_class in self.all_model_classes:
           _lowerCAmelCase   :					Union[str, Any]        =    model_class(A_       )

           _lowerCAmelCase   :					str        =    self._prepare_for_class(A_     ,						A_       )
           _lowerCAmelCase   :					List[str]        =    prepare_numpy_arrays(A_       )

           _lowerCAmelCase   :					Tuple        =    model(A_     ,						noise=A_       )
           _lowerCAmelCase   :					Any        =    model(**A_     ,						noise=A_       )
           self.assert_outputs_same(A_     ,						A_       )

 def    __magic_name__       (		self :	Any     ,						A_ :	List[str]     ,						A_ :	List[str]     ,						A_ :	Optional[Any]       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase   :					Tuple        =    int((tf_model.config.image_size // tf_model.config.patch_size) ** 2       )
      _lowerCAmelCase   :					List[str]        =    np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )
      _lowerCAmelCase   :					Optional[Any]        =    tf.constant(A_       )

      # Add `noise` argument.
      # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
      _lowerCAmelCase   :					List[Any]        =    tf_noise

      super().check_pt_tf_models(A_     ,						A_     ,						A_       )

 def    __magic_name__       (		self :	Union[str, Any]       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase      ,			_lowerCAmelCase   :					Union[str, Any]        =    self.model_tester.prepare_config_and_inputs_for_common()

      _lowerCAmelCase   :					Dict        =    {
          module_member
          for model_class in self.all_model_classes
          for module in (import_module(model_class.__module__       ),)
          for module_member_name in dir(A_       )
          if module_member_name.endswith("MainLayer"       )
          # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
          and module_member_name[: -len("MainLayer"       )] == model_class.__name__[: -len("Model"       )]
          for module_member in (getattr(A_     ,						A_       ),)
          if isinstance(A_     ,						A_       )
          and tf.keras.layers.Layer in module_member.__bases__
          and getattr(A_     ,						"_keras_serializable"     ,						A_       )
      }

      _lowerCAmelCase   :					Union[str, Any]        =    int((config.image_size // config.patch_size) ** 2       )
      _lowerCAmelCase   :					Tuple        =    np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )
      _lowerCAmelCase   :					Union[str, Any]        =    tf.convert_to_tensor(A_       )
      inputs_dict.update({"noise": noise}       )

      for main_layer_class in tf_main_layer_classes:
           _lowerCAmelCase   :					Tuple        =    main_layer_class(A_       )

           _lowerCAmelCase   :					Optional[int]        =    {
               name: tf.keras.Input(tensor.shape[1:]     ,						dtype=tensor.dtype       ) for name, tensor in inputs_dict.items()
           }

           _lowerCAmelCase   :					int        =    tf.keras.Model(A_     ,						outputs=main_layer(A_       )       )
           _lowerCAmelCase   :					Tuple        =    model(A_       )

           with tempfile.TemporaryDirectory() as tmpdirname:
                _lowerCAmelCase   :					List[str]        =    os.path.join(A_     ,						"keras_model.h5"       )
                model.save(A_       )
                _lowerCAmelCase   :					Union[str, Any]        =    tf.keras.models.load_model(
                    A_     ,						custom_objects={main_layer_class.__name__: main_layer_class}       )
                assert isinstance(A_     ,						tf.keras.Model       )
                _lowerCAmelCase   :					Union[str, Any]        =    model(A_       )
                self.assert_outputs_same(A_     ,						A_       )

 @slow
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase      ,			_lowerCAmelCase   :					List[str]        =    self.model_tester.prepare_config_and_inputs_for_common()

      _lowerCAmelCase   :					Optional[int]        =    int((config.image_size // config.patch_size) ** 2       )
      _lowerCAmelCase   :					int        =    np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

      for model_class in self.all_model_classes:
           _lowerCAmelCase   :					Any        =    model_class(A_       )
           _lowerCAmelCase   :					Union[str, Any]        =    self._prepare_for_class(A_     ,						A_       )
           _lowerCAmelCase   :					Any        =    model(A_     ,						noise=A_       )

           if model_class.__name__ == "TFViTMAEModel":
                _lowerCAmelCase   :					List[Any]        =    outputs.last_hidden_state.numpy()
                _lowerCAmelCase   :					str        =    0
           else:
                _lowerCAmelCase   :					Optional[Any]        =    outputs.logits.numpy()
                _lowerCAmelCase   :					Dict        =    0

           with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(A_     ,						saved_model=A_       )
                _lowerCAmelCase   :					Dict        =    model_class.from_pretrained(A_       )
                _lowerCAmelCase   :					Tuple        =    model(A_     ,						noise=A_       )

                if model_class.__name__ == "TFViTMAEModel":
                     _lowerCAmelCase   :					Union[str, Any]        =    after_outputs["last_hidden_state"].numpy()
                     _lowerCAmelCase   :					Optional[Any]        =    0
                else:
                     _lowerCAmelCase   :					Tuple        =    after_outputs["logits"].numpy()
                     _lowerCAmelCase   :					str        =    0

                _lowerCAmelCase   :					Union[str, Any]        =    np.amax(np.abs(out_a - out_a       )       )
                self.assertLessEqual(A_     ,						1E-5       )

 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase      ,			_lowerCAmelCase   :					Optional[Any]        =    self.model_tester.prepare_config_and_inputs_for_common()

      _lowerCAmelCase   :					Union[str, Any]        =    int((config.image_size // config.patch_size) ** 2       )
      _lowerCAmelCase   :					Optional[Any]        =    np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

      for model_class in self.all_model_classes:
           _lowerCAmelCase   :					List[str]        =    model_class(A_       )
           _lowerCAmelCase   :					Optional[Any]        =    self._prepare_for_class(A_     ,						A_       )

           _lowerCAmelCase   :					Union[str, Any]        =    model(A_     ,						noise=A_       )
           _lowerCAmelCase   :					List[str]        =    model.get_config()
           # make sure that returned config is jsonifiable, which is required by keras
           json.dumps(A_       )
           _lowerCAmelCase   :					int        =    model_class.from_config(model.get_config()       )
           # make sure it also accepts a normal config
           _lowerCAmelCase   :					List[str]        =    model_class.from_config(model.config       )
           _lowerCAmelCase   :					Optional[Any]        =    new_model(A_       )  # Build model
           new_model.set_weights(model.get_weights()       )
           _lowerCAmelCase   :					Union[str, Any]        =    new_model(A_     ,						noise=A_       )

           self.assert_outputs_same(A_     ,						A_       )

 @unittest.skip(
     reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."       )
 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      pass

 @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"       )
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      pass


 @slow
 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					str        =    TFViTMAEModel.from_pretrained("google/vit-base-patch16-224"       )
      self.assertIsNotNone(A_       )







def 	_snake_case		(    )       ->      Optional[Any]:

     """simple docstring"""







     _lowerCAmelCase   :					str        =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
     return image





@require_tf
@require_vision
class A__    ( unittest.TestCase ):



 """simple docstring"""







 @cached_property
 def    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"       ) if is_vision_available() else None


 @slow
 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




      np.random.seed(2       )

      _lowerCAmelCase   :					Union[str, Any]        =    TFViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"       )

      _lowerCAmelCase   :					str        =    self.default_image_processor
      _lowerCAmelCase   :					str        =    prepare_img()
      _lowerCAmelCase   :					Optional[int]        =    image_processor(images=A_     ,						return_tensors="tf"       )

      # prepare a noise vector that will be also used for testing the TF model
      # (this way we can ensure that the PT and TF models operate on the same inputs)
      _lowerCAmelCase   :					Dict        =    ViTMAEConfig()
      _lowerCAmelCase   :					str        =    int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2       )
      _lowerCAmelCase   :					Optional[Any]        =    np.random.uniform(size=(1, num_patches)       )

      # forward pass
      _lowerCAmelCase   :					int        =    model(**A_     ,						noise=A_       )

      # verify the logits
      _lowerCAmelCase   :					Optional[int]        =    tf.convert_to_tensor([1, 1_9_6, 7_6_8]       )
      self.assertEqual(outputs.logits.shape     ,						A_       )

      _lowerCAmelCase   :					List[Any]        =    tf.convert_to_tensor(
          [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]       )

      tf.debugging.assert_near(outputs.logits[0, :3, :3]     ,						A_     ,						atol=1E-4       )









from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
  import tensorflow as tf


__UpperCAmelCase    				=     logging.get_logger(__name__)





@dataclass
class A__    ( A ):



 """simple docstring"""







 _lowercase     :					List[Any]      =						[
     '''no_inference''',
     '''no_cuda''',
     '''no_tpu''',
     '''no_speed''',
     '''no_memory''',
     '''no_env_print''',
     '''no_multi_process''',
 ]

 def __init__(		self :	List[Any]     ,						**A_ :	str       ):






      '''simple docstring'''




      for deprecated_arg in self.deprecated_args:
           if deprecated_arg in kwargs:
                _lowerCAmelCase   :					str        =    deprecated_arg[3:]
                _lowerCAmelCase   :					List[Any]        =    not kwargs.pop(A_       )
                logger.warning(
                    F'''{deprecated_arg} is depreciated. Please use --no-{positive_arg} or'''
                    F''' {positive_arg}={kwargs[positive_arg]}'''       )
      _lowerCAmelCase   :					Any        =    kwargs.pop("tpu_name"     ,						self.tpu_name       )
      _lowerCAmelCase   :					str        =    kwargs.pop("device_idx"     ,						self.device_idx       )
      _lowerCAmelCase   :					str        =    kwargs.pop("eager_mode"     ,						self.eager_mode       )
      _lowerCAmelCase   :					Tuple        =    kwargs.pop("use_xla"     ,						self.use_xla       )
      super().__init__(**A_       )

 _lowercase     :					str      =						field(
     default=A  ,		metadata={'''help''': '''Name of TPU'''}  ,		)
 _lowercase     :					int      =						field(
     default=0  ,		metadata={'''help''': '''CPU / GPU device index. Defaults to 0.'''}  ,		)
 _lowercase     :					bool      =						field(default=A  ,		metadata={'''help''': '''Benchmark models in eager model.'''} )
 _lowercase     :					bool      =						field(
     default=A  ,		metadata={
         '''help''': '''Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`.'''
     }  ,		)

 @cached_property
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      _lowerCAmelCase   :					Any        =    None
      if self.tpu:
           try:
                if self.tpu_name:
                     _lowerCAmelCase   :					Union[str, Any]        =    tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name       )
                else:
                     _lowerCAmelCase   :					List[str]        =    tf.distribute.cluster_resolver.TPUClusterResolver()
           except ValueError:
                _lowerCAmelCase   :					Optional[int]        =    None
      return tpu

 @cached_property
 def    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      if self.is_tpu:
           tf.config.experimental_connect_to_cluster(self._setup_tpu       )
           tf.tpu.experimental.initialize_tpu_system(self._setup_tpu       )

           _lowerCAmelCase   :					List[Any]        =    tf.distribute.TPUStrategy(self._setup_tpu       )
      else:
           # currently no multi gpu is allowed
           if self.is_gpu:
                # TODO: Currently only single GPU is supported
                tf.config.set_visible_devices(self.gpu_list[self.device_idx]     ,						"GPU"       )
                _lowerCAmelCase   :					Dict        =    tf.distribute.OneDeviceStrategy(device=F'''/gpu:{self.device_idx}'''       )
           else:
                tf.config.set_visible_devices([]     ,						"GPU"       )  # disable GPU
                _lowerCAmelCase   :					Optional[Any]        =    tf.distribute.OneDeviceStrategy(device=F'''/cpu:{self.device_idx}'''       )

      return strategy

 @property
 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      return self._setup_tpu is not None

 @property
 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      return self._setup_strategy

 @property
 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      return tf.config.list_physical_devices("GPU"       )

 @property
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      if self.cuda:
           return len(self.gpu_list       )
      return 0


 @property
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      return self.n_gpu > 0






