



from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
     import tensorflow as tf

     from transformers import (
         TF_MODEL_FOR_PRETRAINING_MAPPING,
         TFMobileBertForMaskedLM,
         TFMobileBertForMultipleChoice,
         TFMobileBertForNextSentencePrediction,
         TFMobileBertForPreTraining,
         TFMobileBertForQuestionAnswering,
         TFMobileBertForSequenceClassification,
         TFMobileBertForTokenClassification,
         TFMobileBertModel,
     )



@require_tf
class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ,							__SCREAMING_SNAKE_CASE   ,							unittest.TestCase   ):







 '''simple docstring'''
 lowercase_		=      (
     (
         TFMobileBertModel,
         TFMobileBertForMaskedLM,
         TFMobileBertForNextSentencePrediction,
         TFMobileBertForPreTraining,
         TFMobileBertForQuestionAnswering,
         TFMobileBertForSequenceClassification,
         TFMobileBertForTokenClassification,
         TFMobileBertForMultipleChoice,
     )
     if is_tf_available()
     else ()
 )
 lowercase_		=      (
     {
         "feature-extraction": TFMobileBertModel,
         "fill-mask": TFMobileBertForMaskedLM,
         "question-answering": TFMobileBertForQuestionAnswering,
         "text-classification": TFMobileBertForSequenceClassification,
         "token-classification": TFMobileBertForTokenClassification,
         "zero-shot": TFMobileBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 lowercase_		=      False
 lowercase_		=      False

 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Tuple=False)	->Dict:



   '''simple docstring'''



   lowerCamelCase__:				str        =super()._prepare_for_class(UpperCAmelCase_	,			UpperCAmelCase_	,			return_labels=UpperCAmelCase_)

   if return_labels:
     if model_class in get_values(UpperCAmelCase_):
       lowerCamelCase__:				str        =tf.zeros(self.model_tester.batch_size	,			dtype=tf.intaa)

   return inputs_dict



 class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







  '''simple docstring'''

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



    '''simple docstring'''



    lowerCamelCase__:				Optional[Any]        =parent
    lowerCamelCase__:				List[str]        =batch_size
    lowerCamelCase__:				List[str]        =seq_length
    lowerCamelCase__:				List[str]        =is_training
    lowerCamelCase__:				Optional[Any]        =use_input_mask
    lowerCamelCase__:				Any        =use_token_type_ids
    lowerCamelCase__:				List[str]        =use_labels
    lowerCamelCase__:				List[Any]        =vocab_size
    lowerCamelCase__:				List[str]        =hidden_size
    lowerCamelCase__:				int        =num_hidden_layers
    lowerCamelCase__:				str        =num_attention_heads
    lowerCamelCase__:				Dict        =intermediate_size
    lowerCamelCase__:				str        =hidden_act
    lowerCamelCase__:				List[Any]        =hidden_dropout_prob
    lowerCamelCase__:				List[str]        =attention_probs_dropout_prob
    lowerCamelCase__:				Optional[int]        =max_position_embeddings
    lowerCamelCase__:				Union[str, Any]        =type_vocab_size
    lowerCamelCase__:				Any        =type_sequence_label_size
    lowerCamelCase__:				List[str]        =initializer_range
    lowerCamelCase__:				int        =num_labels
    lowerCamelCase__:				Optional[Any]        =num_choices
    lowerCamelCase__:				Any        =scope
    lowerCamelCase__:				Tuple        =embedding_size

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



    '''simple docstring'''



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

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

    lowerCamelCase__:				Optional[Any]        =None
    if self.use_token_type_ids:
      lowerCamelCase__:				Tuple        =ids_tensor([self.batch_size, self.seq_length]	,			self.type_vocab_size)

    lowerCamelCase__:				Optional[Any]        =None
    lowerCamelCase__:				Union[str, Any]        =None
    lowerCamelCase__:				Union[str, Any]        =None
    if self.use_labels:
      lowerCamelCase__:				Dict        =ids_tensor([self.batch_size]	,			self.type_sequence_label_size)
      lowerCamelCase__:				Tuple        =ids_tensor([self.batch_size, self.seq_length]	,			self.num_labels)
      lowerCamelCase__:				Optional[int]        =ids_tensor([self.batch_size]	,			self.num_choices)

    lowerCamelCase__:				Any        =MobileBertConfig(
        vocab_size=self.vocab_size	,			hidden_size=self.hidden_size	,			num_hidden_layers=self.num_hidden_layers	,			num_attention_heads=self.num_attention_heads	,			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	,			max_position_embeddings=self.max_position_embeddings	,			type_vocab_size=self.type_vocab_size	,			initializer_range=self.initializer_range	,			embedding_size=self.embedding_size	,			)

    return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

  def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Tuple	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Tuple)	->Optional[Any]:



    '''simple docstring'''



    lowerCamelCase__:				List[str]        =TFMobileBertModel(config=UpperCAmelCase_)
    lowerCamelCase__:				str        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				Dict        =model(UpperCAmelCase_)

    lowerCamelCase__:				int        =[input_ids, input_mask]
    lowerCamelCase__:				Optional[int]        =model(UpperCAmelCase_)

    lowerCamelCase__:				Tuple        =model(UpperCAmelCase_)

    self.parent.assertEqual(
        result.last_hidden_state.shape	,			(self.batch_size, self.seq_length, self.hidden_size))
    self.parent.assertEqual(result.pooler_output.shape	,			(self.batch_size, self.hidden_size))

  def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Optional[int]	,			UpperCAmelCase_    :       Union[str, Any]	,			UpperCAmelCase_    :       Tuple)	->List[str]:



    '''simple docstring'''



    lowerCamelCase__:				List[str]        =TFMobileBertForMaskedLM(config=UpperCAmelCase_)
    lowerCamelCase__:				Any        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				List[str]        =model(UpperCAmelCase_)
    self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.seq_length, self.vocab_size))

  def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any]	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       Optional[Any])	->Tuple:



    '''simple docstring'''



    lowerCamelCase__:				Optional[int]        =TFMobileBertForNextSentencePrediction(config=UpperCAmelCase_)
    lowerCamelCase__:				List[Any]        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				Union[str, Any]        =model(UpperCAmelCase_)
    self.parent.assertEqual(result.logits.shape	,			(self.batch_size, 2))

  def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       Union[str, Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Union[str, Any])	->List[str]:



    '''simple docstring'''



    lowerCamelCase__:				List[str]        =TFMobileBertForPreTraining(config=UpperCAmelCase_)
    lowerCamelCase__:				List[str]        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				Optional[int]        =model(UpperCAmelCase_)
    self.parent.assertEqual(
        result.prediction_logits.shape	,			(self.batch_size, self.seq_length, self.vocab_size))
    self.parent.assertEqual(result.seq_relationship_logits.shape	,			(self.batch_size, 2))

  def 							SCREAMING_SNAKE_CASE_ (self    :       str	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Union[str, Any]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       List[str])	->Optional[int]:



    '''simple docstring'''



    lowerCamelCase__:				Optional[Any]        =self.num_labels
    lowerCamelCase__:				str        =TFMobileBertForSequenceClassification(config=UpperCAmelCase_)
    lowerCamelCase__:				Tuple        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				Dict        =model(UpperCAmelCase_)
    self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.num_labels))

  def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any]	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       Union[str, Any]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Optional[Any])	->Any:



    '''simple docstring'''



    lowerCamelCase__:				Tuple        =self.num_choices
    lowerCamelCase__:				Union[str, Any]        =TFMobileBertForMultipleChoice(config=UpperCAmelCase_)
    lowerCamelCase__:				str        =tf.tile(tf.expand_dims(UpperCAmelCase_	,			1)	,			(1, self.num_choices, 1))
    lowerCamelCase__:				Union[str, Any]        =tf.tile(tf.expand_dims(UpperCAmelCase_	,			1)	,			(1, self.num_choices, 1))
    lowerCamelCase__:				Any        =tf.tile(tf.expand_dims(UpperCAmelCase_	,			1)	,			(1, self.num_choices, 1))
    lowerCamelCase__:				Tuple        ={
        "input_ids": multiple_choice_inputs_ids,
        "attention_mask": multiple_choice_input_mask,
        "token_type_ids": multiple_choice_token_type_ids,
    }
    lowerCamelCase__:				str        =model(UpperCAmelCase_)
    self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.num_choices))

  def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Union[str, Any]	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[Any])	->List[Any]:



    '''simple docstring'''



    lowerCamelCase__:				Tuple        =self.num_labels
    lowerCamelCase__:				List[Any]        =TFMobileBertForTokenClassification(config=UpperCAmelCase_)
    lowerCamelCase__:				Any        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				str        =model(UpperCAmelCase_)
    self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.seq_length, self.num_labels))

  def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int]	,			UpperCAmelCase_    :       Optional[int]	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[int]	,			UpperCAmelCase_    :       Tuple	,			UpperCAmelCase_    :       Optional[int])	->Optional[int]:



    '''simple docstring'''



    lowerCamelCase__:				List[Any]        =TFMobileBertForQuestionAnswering(config=UpperCAmelCase_)
    lowerCamelCase__:				Optional[Any]        ={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
    lowerCamelCase__:				Union[str, Any]        =model(UpperCAmelCase_)
    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 							SCREAMING_SNAKE_CASE_ (self    :       str)	->Optional[int]:



    '''simple docstring'''



    lowerCamelCase__:				Any        =self.prepare_config_and_inputs()
    (
        lowerCamelCase__ 
    ):				Optional[Any]        =config_and_inputs
    lowerCamelCase__:				Union[str, Any]        ={"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
    return config, inputs_dict



 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Dict:



   '''simple docstring'''



   lowerCamelCase__:				str        =TFMobileBertModelTest.TFMobileBertModelTester(self)
   lowerCamelCase__:				Union[str, Any]        =ConfigTester(self	,			config_class=UpperCAmelCase_	,			hidden_size=37)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->int:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_model(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any)	->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:				str        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_masked_lm(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[int]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_multiple_choice(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       str)	->int:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:				Dict        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_pretraining(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_question_answering(*UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:				str        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_sequence_classification(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->int:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mobilebert_for_token_classification(*UpperCAmelCase_)







 @slow
 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->int:



   '''simple docstring'''



   for model_name in ["google/mobilebert-uncased"]:
     lowerCamelCase__:				Union[str, Any]        =TFMobileBertModel.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)



@require_tf
class        _SCREAMING_SNAKE_CASE						(     unittest.TestCase   ):







 '''simple docstring'''

 @slow
 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->int:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =TFMobileBertForPreTraining.from_pretrained("google/mobilebert-uncased")
   lowerCamelCase__:				str        =tf.constant([[0, 1, 2, 3, 4, 5]])
   lowerCamelCase__:				Tuple        =model(UpperCAmelCase_)[0]

   lowerCamelCase__:				Dict        =[1, 6, 30_522]
   self.assertEqual(output.shape	,			UpperCAmelCase_)

   lowerCamelCase__:				List[str]        =tf.constant(
       [
           [
               [-4.591_9547, -9.24_8295, -9.64_5256],
               [-6.730_6175, -6.44_0284, -6.605_2837],
               [-7.274_3506, -6.784_7915, -6.02_4673],
           ]
       ])
   tf.debugging.assert_near(output[:, :3, :3]	,			UpperCAmelCase_	,			atol=1E-4)








from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''

 def __init__(self    :       int	,			UpperCAmelCase_    :       Distribution	,			UpperCAmelCase_    :       Dict=None	,			UpperCAmelCase_    :       Tuple=None	,			UpperCAmelCase_    :       Tuple=0)	->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =1.0 if scale is None else scale
   lowerCamelCase__:				List[Any]        =0.0 if loc is None else loc

   super().__init__(UpperCAmelCase_	,			[AffineTransform(loc=self.loc	,			scale=self.scale	,			event_dim=UpperCAmelCase_)])

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any])	->int:



   '''simple docstring'''



   return self.base_dist.mean * self.scale + self.loc

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Dict:



   '''simple docstring'''



   return self.base_dist.variance * self.scale**2







 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Optional[Any]:



   '''simple docstring'''



   return self.variance.sqrt()



class        _SCREAMING_SNAKE_CASE						(     nn.Module   ):







 '''simple docstring'''

 def __init__(self    :       Tuple	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Dict[str, int]	,			UpperCAmelCase_    :       Callable[..., Tuple[torch.Tensor]]	,			**UpperCAmelCase_    :       Dict)	->None:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)
   lowerCamelCase__:				Tuple        =args_dim
   lowerCamelCase__:				Any        =nn.ModuleList([nn.Linear(UpperCAmelCase_	,			UpperCAmelCase_) for dim in args_dim.values()])
   lowerCamelCase__:				Any        =domain_map







 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any]	,			UpperCAmelCase_    :       torch.Tensor)	->Tuple[torch.Tensor]:



   '''simple docstring'''



   lowerCamelCase__:				Any        =[proj(UpperCAmelCase_) for proj in self.proj]

   return self.domain_map(*UpperCAmelCase_)



class        _SCREAMING_SNAKE_CASE						(     nn.Module   ):







 '''simple docstring'''

 def __init__(self    :       str	,			UpperCAmelCase_    :       Optional[Any])	->Optional[int]:



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__:				Any        =function







 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			*UpperCAmelCase_    :       Union[str, Any])	->Union[str, Any]:



   '''simple docstring'''



   return self.function(UpperCAmelCase_	,			*UpperCAmelCase_)



class        _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''
 lowercase_		=      42
 lowercase_		=      42
 lowercase_		=      42

 def __init__(self    :       str	,			UpperCAmelCase_    :       int = 1)	->None:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =dim
   lowerCamelCase__:				int        ={k: dim * self.args_dim[k] for k in self.args_dim}

 def 							SCREAMING_SNAKE_CASE_ (self    :       int	,			UpperCAmelCase_    :       str)	->Optional[Any]:



   '''simple docstring'''



   if self.dim == 1:
     return self.distribution_class(*UpperCAmelCase_)
   else:
     return Independent(self.distribution_class(*UpperCAmelCase_)	,			1)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self._base_distribution(UpperCAmelCase_)
   if loc is None and scale is None:
     return distr
   else:
     return AffineTransformed(UpperCAmelCase_	,			loc=UpperCAmelCase_	,			scale=UpperCAmelCase_	,			event_dim=self.event_dim)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Tuple:



   '''simple docstring'''



   return () if self.dim == 1 else (self.dim,)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->int:



   '''simple docstring'''



   return len(self.event_shape)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->float:



   '''simple docstring'''



   return 0.0

 def 							SCREAMING_SNAKE_CASE_ (self    :       str	,			UpperCAmelCase_    :       int)	->nn.Module:



   '''simple docstring'''



   return ParameterProjection(
       in_features=UpperCAmelCase_	,			args_dim=self.args_dim	,			domain_map=LambdaLayer(self.domain_map)	,			)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			*UpperCAmelCase_    :       torch.Tensor)	->Any:



   '''simple docstring'''



   raise NotImplementedError()







 @staticmethod
 def 							SCREAMING_SNAKE_CASE_ (UpperCAmelCase_    :       torch.Tensor)	->torch.Tensor:



   '''simple docstring'''



   return (x + torch.sqrt(torch.square(UpperCAmelCase_) + 4.0)) / 2.0



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"df": 1, "loc": 1, "scale": 1}
 lowercase_		=      StudentT

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[int]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->Dict:



   '''simple docstring'''



   lowerCamelCase__:				str        =cls.squareplus(UpperCAmelCase_).clamp_min(torch.finfo(scale.dtype).eps)
   lowerCamelCase__:				Any        =2.0 + cls.squareplus(UpperCAmelCase_)
   return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"loc": 1, "scale": 1}
 lowercase_		=      Normal

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[Any]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[Any]        =cls.squareplus(UpperCAmelCase_).clamp_min(torch.finfo(scale.dtype).eps)
   return loc.squeeze(-1), scale.squeeze(-1)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"total_count": 1, "logits": 1}
 lowercase_		=      NegativeBinomial

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[int]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				Tuple        =cls.squareplus(UpperCAmelCase_)
   return total_count.squeeze(-1), logits.squeeze(-1)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       Tuple)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Any        =distr_args
   if self.dim == 1:
     return self.distribution_class(total_count=UpperCAmelCase_	,			logits=UpperCAmelCase_)
   else:
     return Independent(self.distribution_class(total_count=UpperCAmelCase_	,			logits=UpperCAmelCase_)	,			1)







 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Union[str, Any]        =distr_args

   if scale is not None:
     # See scaling property of Gamma.
     logits += scale.log()

   return self._base_distribution((total_count, logits))

