


import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
 from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
 from tensorflow.keras.optimizers import Adam




class A  (   tf.keras.optimizers.schedules.LearningRateSchedule     ):






      def __init__(					self						:     str			,   lowercase_						:     float			,   lowercase_						:     Callable			,   lowercase_						:     int			,   lowercase_						:     float = 1.0			,   lowercase_						:     str = None			,   )					->						Optional[int]:
             """simple docstring"""




             super().__init__()
             _lowerCamelCase       :		int		=initial_learning_rate
             _lowerCamelCase       :		Any		=warmup_steps
             _lowerCamelCase       :		Union[str, Any]		=power
             _lowerCamelCase       :		str		=decay_schedule_fn
             _lowerCamelCase       :		List[str]		=name






      def __call__(					self						:     List[Any]			,   lowercase_						:     int )					->						str:
             """simple docstring"""




             with tf.name_scope(self.name or 'WarmUp' ) as name:
                    # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
                    # learning rate will be `global_step/num_warmup_steps * init_lr`.
                    _lowerCamelCase       :		Any		=tf.cast(lowercase_			,   tf.floataa )
                    _lowerCamelCase       :		Dict		=tf.cast(self.warmup_steps			,   tf.floataa )
                    _lowerCamelCase       :		List[str]		=global_step_float / warmup_steps_float
                    _lowerCamelCase       :		int		=self.initial_learning_rate * tf.math.pow(lowercase_			,   self.power )
                    return tf.cond(
                        global_step_float < warmup_steps_float			,   lambda: warmup_learning_rate			,   lambda: self.decay_schedule_fn(step - self.warmup_steps )			,   name=lowercase_			,   )






      def        lowerCamelCase       (					self						:     Any )					->						Any:
             """simple docstring"""




             return {
                 "initial_learning_rate": self.initial_learning_rate,
                 "decay_schedule_fn": self.decay_schedule_fn,
                 "warmup_steps": self.warmup_steps,
                 "power": self.power,
                 "name": self.name,
             }




def 	a_      (    SCREAMING_SNAKE_CASE__					:      float		,			SCREAMING_SNAKE_CASE__					:      int		,			SCREAMING_SNAKE_CASE__					:      int		,			SCREAMING_SNAKE_CASE__					:      float = 0.0		,			SCREAMING_SNAKE_CASE__					:      float = 0.9		,			SCREAMING_SNAKE_CASE__					:      float = 0.9_99		,			SCREAMING_SNAKE_CASE__					:      float = 1e-8		,			SCREAMING_SNAKE_CASE__					:      Optional[float] = None		,			SCREAMING_SNAKE_CASE__					:      Optional[float] = None		,			SCREAMING_SNAKE_CASE__					:      float = 0.0		,			SCREAMING_SNAKE_CASE__					:      float = 1.0		,			SCREAMING_SNAKE_CASE__					:      Optional[List[str]] = None		,			):
       '''simple docstring'''





       _lowerCamelCase       :		Optional[int]		=tf.keras.optimizers.schedules.PolynomialDecay(
           initial_learning_rate=SCREAMING_SNAKE_CASE__		,			decay_steps=num_train_steps - num_warmup_steps		,			end_learning_rate=init_lr * min_lr_ratio		,			power=SCREAMING_SNAKE_CASE__		,			)
       if num_warmup_steps:
              _lowerCamelCase       :		Tuple		=WarmUp(
                  initial_learning_rate=SCREAMING_SNAKE_CASE__		,			decay_schedule_fn=SCREAMING_SNAKE_CASE__		,			warmup_steps=SCREAMING_SNAKE_CASE__		,			)
       if weight_decay_rate > 0.0:
              _lowerCamelCase       :		Optional[int]		=AdamWeightDecay(
                  learning_rate=SCREAMING_SNAKE_CASE__		,			weight_decay_rate=SCREAMING_SNAKE_CASE__		,			beta_a=SCREAMING_SNAKE_CASE__		,			beta_a=SCREAMING_SNAKE_CASE__		,			epsilon=SCREAMING_SNAKE_CASE__		,			clipnorm=SCREAMING_SNAKE_CASE__		,			global_clipnorm=SCREAMING_SNAKE_CASE__		,			exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias']		,			include_in_weight_decay=SCREAMING_SNAKE_CASE__		,			)
       else:
              _lowerCamelCase       :		int		=tf.keras.optimizers.Adam(
                  learning_rate=SCREAMING_SNAKE_CASE__		,			beta_a=SCREAMING_SNAKE_CASE__		,			beta_a=SCREAMING_SNAKE_CASE__		,			epsilon=SCREAMING_SNAKE_CASE__		,			clipnorm=SCREAMING_SNAKE_CASE__		,			global_clipnorm=SCREAMING_SNAKE_CASE__		,			)
       # We return the optimizer and the LR scheduler in order to better track the
       # evolution of the LR independently of the optimizer.
       return optimizer, lr_schedule




class A  (   UpperCamelCase_     ):






      def __init__(					self						:     List[Any]			,   lowercase_						:     Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001			,   lowercase_						:     float = 0.9			,   lowercase_						:     float = 0.999			,   lowercase_						:     float = 1E-7			,   lowercase_						:     bool = False			,   lowercase_						:     float = 0.0			,   lowercase_						:     Optional[List[str]] = None			,   lowercase_						:     Optional[List[str]] = None			,   lowercase_						:     str = "AdamWeightDecay"			,   **lowercase_						:     Dict			,   )					->						List[str]:
             """simple docstring"""




             super().__init__(lowercase_			,   lowercase_			,   lowercase_			,   lowercase_			,   lowercase_			,   lowercase_			,   **lowercase_ )
             _lowerCamelCase       :		Tuple		=weight_decay_rate
             _lowerCamelCase       :		List[Any]		=include_in_weight_decay
             _lowerCamelCase       :		str		=exclude_from_weight_decay






      @classmethod
      def        lowerCamelCase       (					cls						:     List[Any]			,   lowercase_						:     Any )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		int		={'WarmUp': WarmUp}
             return super(lowercase_			,   cls ).from_config(lowercase_			,   custom_objects=lowercase_ )






      def        lowerCamelCase       (					self						:     List[Any]			,   lowercase_						:     str			,   lowercase_						:     Tuple			,   lowercase_						:     Dict )					->						Optional[int]:
             """simple docstring"""




             super(lowercase_			,   self )._prepare_local(lowercase_			,   lowercase_			,   lowercase_ )
             _lowerCamelCase       :		Union[str, Any]		=tf.constant(
                 self.weight_decay_rate			,   name='adam_weight_decay_rate' )






      def        lowerCamelCase       (					self						:     Union[str, Any]			,   lowercase_						:     Optional[int]			,   lowercase_						:     List[str]			,   lowercase_						:     str )					->						Optional[int]:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=self._do_use_weight_decay(var.name )
             if do_decay:
                    return var.assign_sub(
                        learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['weight_decay_rate']			,   use_locking=self._use_locking			,   )
             return tf.no_op()






      def        lowerCamelCase       (					self						:     Union[str, Any]			,   lowercase_						:     Optional[int]			,   lowercase_						:     Union[str, Any]=None			,   **lowercase_						:     int )					->						Any:
             """simple docstring"""




             _lowerCamelCase	,     _lowerCamelCase       :		List[Any]		=list(zip(*lowercase_ ) )
             return super(lowercase_			,   self ).apply_gradients(zip(lowercase_			,   lowercase_ )			,   name=lowercase_			,   **lowercase_ )






      def        lowerCamelCase       (					self						:     Union[str, Any]			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Optional[int]			,   lowercase_						:     str )					->						List[Any]:
             """simple docstring"""




             if apply_state is None:
                    return self._decayed_lr_t[var_dtype], {}

             _lowerCamelCase       :		int		=apply_state or {}
             _lowerCamelCase       :		List[str]		=apply_state.get((var_device, var_dtype) )
             if coefficients is None:
                    _lowerCamelCase       :		List[str]		=self._fallback_apply_state(lowercase_			,   lowercase_ )
                    _lowerCamelCase       :		Union[str, Any]		=coefficients

             return coefficients["lr_t"], {"apply_state": apply_state}






      def        lowerCamelCase       (					self						:     Tuple			,   lowercase_						:     Any			,   lowercase_						:     List[Any]			,   lowercase_						:     List[Any]=None )					->						Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase	,     _lowerCamelCase       :		Optional[int]		=self._get_lr(var.device			,   var.dtype.base_dtype			,   lowercase_ )
             _lowerCamelCase       :		Tuple		=self._decay_weights_op(lowercase_			,   lowercase_			,   lowercase_ )
             with tf.control_dependencies([decay] ):
                    return super(lowercase_			,   self )._resource_apply_dense(lowercase_			,   lowercase_			,   **lowercase_ )






      def        lowerCamelCase       (					self						:     int			,   lowercase_						:     str			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Dict			,   lowercase_						:     List[str]=None )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase	,     _lowerCamelCase       :		int		=self._get_lr(var.device			,   var.dtype.base_dtype			,   lowercase_ )
             _lowerCamelCase       :		Dict		=self._decay_weights_op(lowercase_			,   lowercase_			,   lowercase_ )
             with tf.control_dependencies([decay] ):
                    return super(lowercase_			,   self )._resource_apply_sparse(lowercase_			,   lowercase_			,   lowercase_			,   **lowercase_ )






      def        lowerCamelCase       (					self						:     Dict )					->						Optional[int]:
             """simple docstring"""




             _lowerCamelCase       :		Tuple		=super().get_config()
             config.update({'weight_decay_rate': self.weight_decay_rate} )
             return config






      def        lowerCamelCase       (					self						:     List[Any]			,   lowercase_						:     Dict )					->						List[str]:
             """simple docstring"""




             if self.weight_decay_rate == 0:
                    return False

             if self._include_in_weight_decay:
                    for r in self._include_in_weight_decay:
                           if re.search(lowercase_			,   lowercase_ ) is not None:
                                  return True

             if self._exclude_from_weight_decay:
                    for r in self._exclude_from_weight_decay:
                           if re.search(lowercase_			,   lowercase_ ) is not None:
                                  return False
             return True




class A  (   UpperCamelCase_     ):






      def __init__(					self						:     Union[str, Any] )					->						List[str]:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=[]
             _lowerCamelCase       :		Tuple		=None






      @property
      def        lowerCamelCase       (					self						:     Union[str, Any] )					->						str:
             """simple docstring"""




             if self._accum_steps is None:
                    _lowerCamelCase       :		Any		=tf.Variable(
                        tf.constant(0			,   dtype=tf.intaa )			,   trainable=lowercase_			,   synchronization=tf.VariableSynchronization.ON_READ			,   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA			,   )

             return self._accum_steps.value()






      @property
      def        lowerCamelCase       (					self						:     Tuple )					->						Dict:
             """simple docstring"""




             if not self._gradients:
                    raise ValueError('The accumulator should be called first to initialize the gradients' )
             return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]






      def __call__(					self						:     int			,   lowercase_						:     int )					->						Tuple:
             """simple docstring"""




             if not self._gradients:
                    _lowerCamelCase       :		List[str]		=self.step  # Create the step variable.
                    self._gradients.extend(
                        [
                            tf.Variable(
                                tf.zeros_like(lowercase_ )			,   trainable=lowercase_			,   synchronization=tf.VariableSynchronization.ON_READ			,   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA			,   )
                            if gradient is not None
                            else gradient
                            for gradient in gradients
                        ] )
             if len(lowercase_ ) != len(self._gradients ):
                    raise ValueError(F'''Expected {len(self._gradients )} gradients, but got {len(lowercase_ )}''' )

             for accum_gradient, gradient in zip(self._gradients			,   lowercase_ ):
                    if accum_gradient is not None and gradient is not None:
                           accum_gradient.assign_add(lowercase_ )

             self._accum_steps.assign_add(1 )






      def        lowerCamelCase       (					self						:     str )					->						int:
             """simple docstring"""




             if not self._gradients:
                    return
             self._accum_steps.assign(0 )
             for gradient in self._gradients:
                    if gradient is not None:
                           gradient.assign(tf.zeros_like(lowercase_ ) )






import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     Any )					->						Any:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]
             _lowerCamelCase       :		Dict		=[]

             for i in range(self.num_layers ):
                    _lowerCamelCase       :		Union[str, Any]		=self.in_channels if i == 0 else self.out_channels

                    _lowerCamelCase       :		Any		=FlaxResnetBlockaD(
                        in_channels=lowercase_			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

                    _lowerCamelCase       :		Dict		=FlaxTransformeraDModel(
                        in_channels=self.out_channels			,   n_heads=self.num_attention_heads			,   d_head=self.out_channels // self.num_attention_heads			,   depth=1			,   use_linear_projection=self.use_linear_projection			,   only_cross_attention=self.only_cross_attention			,   use_memory_efficient_attention=self.use_memory_efficient_attention			,   dtype=self.dtype			,   )
                    attentions.append(lowercase_ )

             _lowerCamelCase       :		Optional[int]		=resnets
             _lowerCamelCase       :		Dict		=attentions

             if self.add_downsample:
                    _lowerCamelCase       :		Union[str, Any]		=FlaxDownsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     Any			,   lowercase_						:     Optional[Any]			,   lowercase_						:     List[Any]			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     List[str]=True )					->						Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase       :		Tuple		=()

             for resnet, attn in zip(self.resnets			,   self.attentions ):
                    _lowerCamelCase       :		Union[str, Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		Union[str, Any]		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    output_states += (hidden_states,)

             if self.add_downsample:
                    _lowerCamelCase       :		Optional[int]		=self.downsamplers_a(lowercase_ )
                    output_states += (hidden_states,)

             return hidden_states, output_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     Optional[int] )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]

             for i in range(self.num_layers ):
                    _lowerCamelCase       :		Tuple		=self.in_channels if i == 0 else self.out_channels

                    _lowerCamelCase       :		Any		=FlaxResnetBlockaD(
                        in_channels=lowercase_			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )
             _lowerCamelCase       :		Union[str, Any]		=resnets

             if self.add_downsample:
                    _lowerCamelCase       :		List[str]		=FlaxDownsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     Any			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Optional[Any]=True )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		Optional[int]		=()

             for resnet in self.resnets:
                    _lowerCamelCase       :		Tuple		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    output_states += (hidden_states,)

             if self.add_downsample:
                    _lowerCamelCase       :		Tuple		=self.downsamplers_a(lowercase_ )
                    output_states += (hidden_states,)

             return hidden_states, output_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     Dict )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]
             _lowerCamelCase       :		List[str]		=[]

             for i in range(self.num_layers ):
                    _lowerCamelCase       :		List[str]		=self.in_channels if (i == self.num_layers - 1) else self.out_channels
                    _lowerCamelCase       :		Tuple		=self.prev_output_channel if i == 0 else self.out_channels

                    _lowerCamelCase       :		List[str]		=FlaxResnetBlockaD(
                        in_channels=resnet_in_channels + res_skip_channels			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

                    _lowerCamelCase       :		Tuple		=FlaxTransformeraDModel(
                        in_channels=self.out_channels			,   n_heads=self.num_attention_heads			,   d_head=self.out_channels // self.num_attention_heads			,   depth=1			,   use_linear_projection=self.use_linear_projection			,   only_cross_attention=self.only_cross_attention			,   use_memory_efficient_attention=self.use_memory_efficient_attention			,   dtype=self.dtype			,   )
                    attentions.append(lowercase_ )

             _lowerCamelCase       :		int		=resnets
             _lowerCamelCase       :		Dict		=attentions

             if self.add_upsample:
                    _lowerCamelCase       :		str		=FlaxUpsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     Optional[Any]			,   lowercase_						:     Any			,   lowercase_						:     Any			,   lowercase_						:     List[Any]			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any]=True )					->						Optional[int]:
             """simple docstring"""




             for resnet, attn in zip(self.resnets			,   self.attentions ):
                    # pop res hidden states
                    _lowerCamelCase       :		Optional[int]		=res_hidden_states_tuple[-1]
                    _lowerCamelCase       :		Union[str, Any]		=res_hidden_states_tuple[:-1]
                    _lowerCamelCase       :		Any		=jnp.concatenate((hidden_states, res_hidden_states)			,   axis=-1 )

                    _lowerCamelCase       :		Optional[Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		List[Any]		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             if self.add_upsample:
                    _lowerCamelCase       :		Optional[Any]		=self.upsamplers_a(lowercase_ )

             return hidden_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     List[Any] )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=[]

             for i in range(self.num_layers ):
                    _lowerCamelCase       :		Tuple		=self.in_channels if (i == self.num_layers - 1) else self.out_channels
                    _lowerCamelCase       :		int		=self.prev_output_channel if i == 0 else self.out_channels

                    _lowerCamelCase       :		str		=FlaxResnetBlockaD(
                        in_channels=resnet_in_channels + res_skip_channels			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

             _lowerCamelCase       :		str		=resnets

             if self.add_upsample:
                    _lowerCamelCase       :		List[str]		=FlaxUpsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     List[str]			,   lowercase_						:     str			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Any=True )					->						int:
             """simple docstring"""




             for resnet in self.resnets:
                    # pop res hidden states
                    _lowerCamelCase       :		List[str]		=res_hidden_states_tuple[-1]
                    _lowerCamelCase       :		str		=res_hidden_states_tuple[:-1]
                    _lowerCamelCase       :		List[str]		=jnp.concatenate((hidden_states, res_hidden_states)			,   axis=-1 )

                    _lowerCamelCase       :		Optional[Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             if self.add_upsample:
                    _lowerCamelCase       :		Union[str, Any]		=self.upsamplers_a(lowercase_ )

             return hidden_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     int )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=[
                 FlaxResnetBlockaD(
                     in_channels=self.in_channels			,   out_channels=self.in_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
             ]

             _lowerCamelCase       :		Any		=[]

             for _ in range(self.num_layers ):
                    _lowerCamelCase       :		Optional[int]		=FlaxTransformeraDModel(
                        in_channels=self.in_channels			,   n_heads=self.num_attention_heads			,   d_head=self.in_channels // self.num_attention_heads			,   depth=1			,   use_linear_projection=self.use_linear_projection			,   use_memory_efficient_attention=self.use_memory_efficient_attention			,   dtype=self.dtype			,   )
                    attentions.append(lowercase_ )

                    _lowerCamelCase       :		Tuple		=FlaxResnetBlockaD(
                        in_channels=self.in_channels			,   out_channels=self.in_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

             _lowerCamelCase       :		List[Any]		=resnets
             _lowerCamelCase       :		List[str]		=attentions






      def __call__(					self						:     Optional[int]			,   lowercase_						:     List[Any]			,   lowercase_						:     Dict			,   lowercase_						:     Tuple			,   lowercase_						:     List[str]=True )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		Dict		=self.resnets[0](lowercase_			,   lowercase_ )
             for attn, resnet in zip(self.attentions			,   self.resnets[1:] ):
                    _lowerCamelCase       :		Tuple		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		List[str]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             return hidden_states


