

import os
import tempfile
import unittest

from transformers import DistilBertConfig, 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 (
			    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
			    DistilBertForMaskedLM,
			    DistilBertForMultipleChoice,
			    DistilBertForQuestionAnswering,
			    DistilBertForSequenceClassification,
			    DistilBertForTokenClassification,
			    DistilBertModel,
			)






class _UpperCamelCase							(							__A      ):
			'''simple docstring'''







			def __init__(	self							:       List[str]					,					a							:       List[Any]					,					a							:       List[Any]=13					,					a							:       Dict=7					,					a							:       Optional[Any]=True					,					a							:       List[Any]=True					,					a							:       Tuple=False					,					a							:       Dict=True					,					a							:       Dict=99					,					a							:       List[str]=32					,					a							:       Dict=5					,					a							:       Dict=4					,					a							:       Tuple=37					,					a							:       str="gelu"					,					a							:       Optional[Any]=0.1					,					a							:       Optional[int]=0.1					,					a							:       int=512					,					a							:       Union[str, Any]=16					,					a							:       Dict=2					,					a							:       List[str]=0.02					,					a							:       Any=3					,					a							:       Tuple=4					,					a							:       Union[str, Any]=None					,					)							->				Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				parent
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				batch_size
				SCREAMING_SNAKE_CASE      :					List[str]    =				seq_length
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				is_training
				SCREAMING_SNAKE_CASE      :					str    =				use_input_mask
				SCREAMING_SNAKE_CASE      :					List[str]    =				use_token_type_ids
				SCREAMING_SNAKE_CASE      :					int    =				use_labels
				SCREAMING_SNAKE_CASE      :					Any    =				vocab_size
				SCREAMING_SNAKE_CASE      :					List[str]    =				hidden_size
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				num_hidden_layers
				SCREAMING_SNAKE_CASE      :					Any    =				num_attention_heads
				SCREAMING_SNAKE_CASE      :					List[Any]    =				intermediate_size
				SCREAMING_SNAKE_CASE      :					List[Any]    =				hidden_act
				SCREAMING_SNAKE_CASE      :					Any    =				hidden_dropout_prob
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE      :					Dict    =				max_position_embeddings
				SCREAMING_SNAKE_CASE      :					str    =				type_vocab_size
				SCREAMING_SNAKE_CASE      :					str    =				type_sequence_label_size
				SCREAMING_SNAKE_CASE      :					Any    =				initializer_range
				SCREAMING_SNAKE_CASE      :					Any    =				num_labels
				SCREAMING_SNAKE_CASE      :					List[Any]    =				num_choices
				SCREAMING_SNAKE_CASE      :					Dict    =				scope
			def 							__UpperCamelCase       (	self							:       str		)							->				Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Dict    =				ids_tensor([self.batch_size, self.seq_length]					,					self.vocab_size		)

				SCREAMING_SNAKE_CASE      :					List[str]    =				None
				if self.use_input_mask:
					SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				random_attention_mask([self.batch_size, self.seq_length]		)

				SCREAMING_SNAKE_CASE      :					Tuple    =				None
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				None
				SCREAMING_SNAKE_CASE      :					Tuple    =				None
				if self.use_labels:
					SCREAMING_SNAKE_CASE      :					Dict    =				ids_tensor([self.batch_size]					,					self.type_sequence_label_size		)
					SCREAMING_SNAKE_CASE      :					Dict    =				ids_tensor([self.batch_size, self.seq_length]					,					self.num_labels		)
					SCREAMING_SNAKE_CASE      :					List[Any]    =				ids_tensor([self.batch_size]					,					self.num_choices		)

				SCREAMING_SNAKE_CASE      :					List[Any]    =				self.get_config()

				return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
			def 							__UpperCamelCase       (	self							:       Union[str, Any]		)							->				int:

				"""simple docstring"""

				return DistilBertConfig(
				    vocab_size=self.vocab_size					,					dim=self.hidden_size					,					n_layers=self.num_hidden_layers					,					n_heads=self.num_attention_heads					,					hidden_dim=self.intermediate_size					,					hidden_act=self.hidden_act					,					dropout=self.hidden_dropout_prob					,					attention_dropout=self.attention_probs_dropout_prob					,					max_position_embeddings=self.max_position_embeddings					,					initializer_range=self.initializer_range					,					)
			def 							__UpperCamelCase       (	self							:       str					,					a							:       Any					,					a							:       Optional[Any]					,					a							:       List[str]					,					a							:       List[str]					,					a							:       Union[str, Any]					,					a							:       List[str]		)							->				List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				DistilBertModel(config=a		)
				model.to(a		)
				model.eval()
				SCREAMING_SNAKE_CASE      :					Tuple    =				model(a					,					a		)
				SCREAMING_SNAKE_CASE      :					str    =				model(a		)
				self.parent.assertEqual(result.last_hidden_state.shape					,					(self.batch_size, self.seq_length, self.hidden_size)		)
			def 							__UpperCamelCase       (	self							:       Optional[Any]					,					a							:       Dict					,					a							:       Tuple					,					a							:       Optional[int]					,					a							:       Optional[Any]					,					a							:       Any					,					a							:       int		)							->				List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Optional[int]    =				DistilBertForMaskedLM(config=a		)
				model.to(a		)
				model.eval()
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				model(a					,					attention_mask=a					,					labels=a		)
				self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.seq_length, self.vocab_size)		)
			def 							__UpperCamelCase       (	self							:       str					,					a							:       Optional[Any]					,					a							:       int					,					a							:       Tuple					,					a							:       Union[str, Any]					,					a							:       int					,					a							:       Optional[Any]		)							->				Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[Any]    =				DistilBertForQuestionAnswering(config=a		)
				model.to(a		)
				model.eval()
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				model(
				    a					,					attention_mask=a					,					start_positions=a					,					end_positions=a		)
				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 							__UpperCamelCase       (	self							:       Tuple					,					a							:       str					,					a							:       Any					,					a							:       Optional[int]					,					a							:       Dict					,					a							:       Dict					,					a							:       Union[str, Any]		)							->				Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Any    =				self.num_labels
				SCREAMING_SNAKE_CASE      :					List[Any]    =				DistilBertForSequenceClassification(a		)
				model.to(a		)
				model.eval()
				SCREAMING_SNAKE_CASE      :					List[str]    =				model(a					,					attention_mask=a					,					labels=a		)
				self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.num_labels)		)
			def 							__UpperCamelCase       (	self							:       Optional[int]					,					a							:       Dict					,					a							:       Any					,					a							:       List[Any]					,					a							:       List[Any]					,					a							:       Dict					,					a							:       List[str]		)							->				Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					str    =				self.num_labels
				SCREAMING_SNAKE_CASE      :					Optional[int]    =				DistilBertForTokenClassification(config=a		)
				model.to(a		)
				model.eval()

				SCREAMING_SNAKE_CASE      :					Dict    =				model(a					,					attention_mask=a					,					labels=a		)
				self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.seq_length, self.num_labels)		)
			def 							__UpperCamelCase       (	self							:       str					,					a							:       Any					,					a							:       int					,					a							:       Dict					,					a							:       int					,					a							:       Optional[Any]					,					a							:       List[str]		)							->				Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[Any]    =				self.num_choices
				SCREAMING_SNAKE_CASE      :					Optional[int]    =				DistilBertForMultipleChoice(config=a		)
				model.to(a		)
				model.eval()
				SCREAMING_SNAKE_CASE      :					Tuple    =				input_ids.unsqueeze(1		).expand(-1					,					self.num_choices					,					-1		).contiguous()
				SCREAMING_SNAKE_CASE      :					str    =				input_mask.unsqueeze(1		).expand(-1					,					self.num_choices					,					-1		).contiguous()
				SCREAMING_SNAKE_CASE      :					Tuple    =				model(
				    a					,					attention_mask=a					,					labels=a					,					)
				self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.num_choices)		)






			def 							__UpperCamelCase       (	self							:       Optional[Any]		)							->				List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Dict    =				self.prepare_config_and_inputs()
				((SCREAMING_SNAKE_CASE)   ,(SCREAMING_SNAKE_CASE)   ,(SCREAMING_SNAKE_CASE)   ,(SCREAMING_SNAKE_CASE)   ,(SCREAMING_SNAKE_CASE)   ,(SCREAMING_SNAKE_CASE))      :					Optional[Any]    =				config_and_inputs
				SCREAMING_SNAKE_CASE      :					List[Any]    =				{"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict






@require_torch
class _UpperCamelCase							(							__A							,		__A							,		unittest.TestCase      ):
			'''simple docstring'''







			lowerCamelCase__							=(
			    (
			        DistilBertModel,
			        DistilBertForMaskedLM,
			        DistilBertForMultipleChoice,
			        DistilBertForQuestionAnswering,
			        DistilBertForSequenceClassification,
			        DistilBertForTokenClassification,
			    )
			    if is_torch_available()
			    else None
			)
			lowerCamelCase__							=(
			    {
			        'feature-extraction': DistilBertModel,
			        'fill-mask': DistilBertForMaskedLM,
			        'question-answering': DistilBertForQuestionAnswering,
			        'text-classification': DistilBertForSequenceClassification,
			        'token-classification': DistilBertForTokenClassification,
			        'zero-shot': DistilBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__							=True
			lowerCamelCase__							=True
			lowerCamelCase__							=True
			lowerCamelCase__							=True
			def 							__UpperCamelCase       (	self							:       Dict		)							->				Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[Any]    =				DistilBertModelTester(self		)
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				ConfigTester(self					,					config_class=a					,					dim=37		)
			def 							__UpperCamelCase       (	self							:       List[str]		)							->				Dict:

				"""simple docstring"""

				self.config_tester.run_common_tests()
			def 							__UpperCamelCase       (	self							:       str		)							->				Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Any    =				self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_model(*a		)
			def 							__UpperCamelCase       (	self							:       Any		)							->				List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					str    =				self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_masked_lm(*a		)
			def 							__UpperCamelCase       (	self							:       Any		)							->				Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_question_answering(*a		)
			def 							__UpperCamelCase       (	self							:       str		)							->				str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Optional[int]    =				self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_sequence_classification(*a		)
			def 							__UpperCamelCase       (	self							:       Union[str, Any]		)							->				Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[str]    =				self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_token_classification(*a		)
			def 							__UpperCamelCase       (	self							:       List[str]		)							->				List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					str    =				self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_multiple_choice(*a		)
			@slow
			def 							__UpperCamelCase       (	self							:       Any		)							->				Tuple:

				"""simple docstring"""

				for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					SCREAMING_SNAKE_CASE      :					Dict    =				DistilBertModel.from_pretrained(a		)
					self.assertIsNotNone(a		)






			@slow
			@require_torch_gpu
			def 							__UpperCamelCase       (	self							:       Union[str, Any]		)							->				Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   ,SCREAMING_SNAKE_CASE      :					Optional[Any]    =				self.model_tester.prepare_config_and_inputs_for_common()
				for model_class in self.all_model_classes:
					# BertForMultipleChoice behaves incorrectly in JIT environments.
					if model_class == DistilBertForMultipleChoice:
						return

					SCREAMING_SNAKE_CASE      :					int    =				True
					SCREAMING_SNAKE_CASE      :					Any    =				model_class(config=a		)

					SCREAMING_SNAKE_CASE      :					int    =				self._prepare_for_class(a					,					a		)
					SCREAMING_SNAKE_CASE      :					Dict    =				torch.jit.trace(
					    a					,					(inputs_dict["input_ids"].to("cpu"		), inputs_dict["attention_mask"].to("cpu"		))		)

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






@require_torch
class _UpperCamelCase							(							unittest.TestCase      ):
			'''simple docstring'''







			@slow
			def 							__UpperCamelCase       (	self							:       Optional[int]		)							->				Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Tuple    =				DistilBertModel.from_pretrained("distilbert-base-uncased"		)
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
				SCREAMING_SNAKE_CASE      :					int    =				torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
				with torch.no_grad():
					SCREAMING_SNAKE_CASE      :					Optional[Any]    =				model(a					,					attention_mask=a		)[0]
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				torch.Size((1, 11, 768)		)
				self.assertEqual(output.shape					,					a		)
				SCREAMING_SNAKE_CASE      :					Any    =				torch.tensor(
				    [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]		)

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




"""simple docstring"""







from ..utils import DummyObject, requires_backends







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      List[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Any    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      int    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      Dict    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      int    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      str       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Dict    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      List[str]    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[Any]    ,			*SCREAMING_SNAKE_CASE_				:      str    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[str]    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Any    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      str    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      int    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Any    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      str       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			**SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      str       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Optional[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Dict    ,			**SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      str    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )





