code: 
  
  
  
  
  
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  
  
  
  def    _lowerCamelCase	(	__lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase = 4096					)   -> Union[str, Any]:
  
       '''simple docstring'''
       UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			:     Optional[Any]										= key.shape[-3:]
       UpperCAmelCase__			:     List[Any]										= value.shape[-1]
       UpperCAmelCase__			:     str										= min(__lowerCamelCase   ,       __lowerCamelCase					)
       UpperCAmelCase__			:     str										= query / jnp.sqrt(__lowerCamelCase					)
  
       @functools.partial(jax.checkpoint   ,       prevent_cse=__lowerCamelCase					)
       def summarize_chunk(__lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase					):
            UpperCAmelCase__			:     Optional[int]										= jnp.einsum("""...qhd,...khd->...qhk"""   ,       __lowerCamelCase   ,       __lowerCamelCase   ,       precision=__lowerCamelCase					)
  
            UpperCAmelCase__			:     Optional[int]										= jnp.max(__lowerCamelCase   ,       axis=-1   ,       keepdims=__lowerCamelCase					)
            UpperCAmelCase__			:     Union[str, Any]										= jax.lax.stop_gradient(__lowerCamelCase					)
            UpperCAmelCase__			:     str										= jnp.exp(attn_weights - max_score					)
  
            UpperCAmelCase__			:     Any										= jnp.einsum("""...vhf,...qhv->...qhf"""   ,       __lowerCamelCase   ,       __lowerCamelCase   ,       precision=__lowerCamelCase					)
            UpperCAmelCase__			:     Optional[int]										= jnp.einsum("""...qhk->...qh"""   ,       __lowerCamelCase					)
  
            return (exp_values, exp_weights.sum(axis=-1					), max_score)
  
       def chunk_scanner(__lowerCamelCase					):
            # julienne key array
            UpperCAmelCase__			:     Union[str, Any]										= jax.lax.dynamic_slice(
                operand=__lowerCamelCase   ,       start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]   ,       slice_sizes=list(key.shape[:-3]					) + [key_chunk_size, num_heads, k_features]   ,       )
  
            # julienne value array
            UpperCAmelCase__			:     Dict										= jax.lax.dynamic_slice(
                operand=__lowerCamelCase   ,       start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]   ,       slice_sizes=list(value.shape[:-3]					) + [key_chunk_size, num_heads, v_features]   ,       )
  
            return summarize_chunk(__lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase					)
  
       UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			:     int										= jax.lax.map(f=__lowerCamelCase   ,       xs=jnp.arange(0   ,       __lowerCamelCase   ,       __lowerCamelCase					)					)
  
       UpperCAmelCase__			:     Optional[Any]										= jnp.max(__lowerCamelCase   ,       axis=0   ,       keepdims=__lowerCamelCase					)
       UpperCAmelCase__			:     Dict										= jnp.exp(chunk_max - global_max					)
  
       chunk_values *= jnp.expand_dims(__lowerCamelCase   ,       axis=-1					)
       chunk_weights *= max_diffs
  
       UpperCAmelCase__			:     Tuple										= chunk_values.sum(axis=0					)
       UpperCAmelCase__			:     List[str]										= jnp.expand_dims(__lowerCamelCase   ,       -1					).sum(axis=0					)
  
       return all_values / all_weights
  
  
  
  
  
  def    _lowerCamelCase	(	__lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase=jax.lax.Precision.HIGHEST   ,       __lowerCamelCase = 1024   ,       __lowerCamelCase = 4096					)   -> int:
  
       '''simple docstring'''
       UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			:     str										= query.shape[-3:]
  
       def chunk_scanner(__lowerCamelCase   ,       __lowerCamelCase					):
            # julienne query array
            UpperCAmelCase__			:     Optional[int]										= jax.lax.dynamic_slice(
                operand=__lowerCamelCase   ,       start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]   ,       slice_sizes=list(query.shape[:-3]					) + [min(__lowerCamelCase   ,       __lowerCamelCase					), num_heads, q_features]   ,       )
  
            return (
                chunk_idx + query_chunk_size,  # unused ignore it
                _query_chunk_attention(
                    query=__lowerCamelCase   ,       key=__lowerCamelCase   ,       value=__lowerCamelCase   ,       precision=__lowerCamelCase   ,       key_chunk_size=__lowerCamelCase					),
            )
  
       UpperCAmelCase__			,		UpperCAmelCase__			:     int										= jax.lax.scan(
           f=__lowerCamelCase   ,       init=0   ,       xs=__lowerCamelCase   ,       length=math.ceil(num_q / query_chunk_size					)					)
  
       return jnp.concatenate(__lowerCamelCase   ,       axis=-3					)  # fuse the chunked result back
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     (		nn.Module		):
       __lowerCamelCase										=     42
       __lowerCamelCase										=     8
       __lowerCamelCase										=     64
       __lowerCamelCase										=     0.0
       __lowerCamelCase										=     False
       __lowerCamelCase										=     jnp.floataa
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            UpperCAmelCase__			:     str										= self.dim_head * self.heads
            UpperCAmelCase__			:     Optional[Any]										= self.dim_head**-0.5
  
            # Weights were exported with old names {to_q, to_k, to_v, to_out}
            UpperCAmelCase__			:     Any										= nn.Dense(_lowerCAmelCase  ,	use_bias=_lowerCAmelCase  ,	dtype=self.dtype  ,	name="""to_q"""		)
            UpperCAmelCase__			:     int										= nn.Dense(_lowerCAmelCase  ,	use_bias=_lowerCAmelCase  ,	dtype=self.dtype  ,	name="""to_k"""		)
            UpperCAmelCase__			:     Union[str, Any]										= nn.Dense(_lowerCAmelCase  ,	use_bias=_lowerCAmelCase  ,	dtype=self.dtype  ,	name="""to_v"""		)
  
            UpperCAmelCase__			:     int										= nn.Dense(self.query_dim  ,	dtype=self.dtype  ,	name="""to_out_0"""		)
            UpperCAmelCase__			:     Union[str, Any]										= nn.Dropout(rate=self.dropout		)
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase		):
            UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			:     Optional[Any]										= tensor.shape
            UpperCAmelCase__			:     Tuple										= self.heads
            UpperCAmelCase__			:     Tuple										= tensor.reshape(_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	dim // head_size		)
            UpperCAmelCase__			:     Tuple										= jnp.transpose(_lowerCAmelCase  ,	(0, 2, 1, 3)		)
            UpperCAmelCase__			:     str										= tensor.reshape(batch_size * head_size  ,	_lowerCAmelCase  ,	dim // head_size		)
            return tensor
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase		):
            UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			:     Union[str, Any]										= tensor.shape
            UpperCAmelCase__			:     List[Any]										= self.heads
            UpperCAmelCase__			:     Optional[Any]										= tensor.reshape(batch_size // head_size  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase		)
            UpperCAmelCase__			:     List[Any]										= jnp.transpose(_lowerCAmelCase  ,	(0, 2, 1, 3)		)
            UpperCAmelCase__			:     List[Any]										= tensor.reshape(batch_size // head_size  ,	_lowerCAmelCase  ,	dim * head_size		)
            return tensor
  
  
  
  
       def __call__(							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase=None  ,	_lowerCAmelCase=True		):
            UpperCAmelCase__			:     Union[str, Any]										= hidden_states if context is None else context
  
            UpperCAmelCase__			:     int										= self.query(_lowerCAmelCase		)
            UpperCAmelCase__			:     List[str]										= self.key(_lowerCAmelCase		)
            UpperCAmelCase__			:     Union[str, Any]										= self.value(_lowerCAmelCase		)
  
            UpperCAmelCase__			:     Union[str, Any]										= self.reshape_heads_to_batch_dim(_lowerCAmelCase		)
            UpperCAmelCase__			:     Union[str, Any]										= self.reshape_heads_to_batch_dim(_lowerCAmelCase		)
            UpperCAmelCase__			:     Dict										= self.reshape_heads_to_batch_dim(_lowerCAmelCase		)
  
            if self.use_memory_efficient_attention:
                 UpperCAmelCase__			:     List[str]										= query_states.transpose(1  ,	0  ,	2		)
                 UpperCAmelCase__			:     Tuple										= key_states.transpose(1  ,	0  ,	2		)
                 UpperCAmelCase__			:     List[Any]										= value_states.transpose(1  ,	0  ,	2		)
  
                 # this if statement create a chunk size for each layer of the unet
                 # the chunk size is equal to the query_length dimension of the deepest layer of the unet
  
                 UpperCAmelCase__			:     List[str]										= query_states.shape[-3]
                 if flatten_latent_dim % 64 == 0:
                      UpperCAmelCase__			:     Dict										= int(flatten_latent_dim / 64		)
                 elif flatten_latent_dim % 16 == 0:
                      UpperCAmelCase__			:     int										= int(flatten_latent_dim / 16		)
                 elif flatten_latent_dim % 4 == 0:
                      UpperCAmelCase__			:     str										= int(flatten_latent_dim / 4		)
                 else:
                      UpperCAmelCase__			:     Tuple										= int(_lowerCAmelCase		)
  
                 UpperCAmelCase__			:     int										= jax_memory_efficient_attention(
                     _lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	query_chunk_size=_lowerCAmelCase  ,	key_chunk_size=4096 * 4		)
  
                 UpperCAmelCase__			:     List[Any]										= hidden_states.transpose(1  ,	0  ,	2		)
            else:
                 # compute attentions
                 UpperCAmelCase__			:     Tuple										= jnp.einsum("""b i d, b j d->b i j"""  ,	_lowerCAmelCase  ,	_lowerCAmelCase		)
                 UpperCAmelCase__			:     Any										= attention_scores * self.scale
                 UpperCAmelCase__			:     Optional[int]										= nn.softmax(_lowerCAmelCase  ,	axis=2		)
  
                 # attend to values
                 UpperCAmelCase__			:     Any										= jnp.einsum("""b i j, b j d -> b i d"""  ,	_lowerCAmelCase  ,	_lowerCAmelCase		)
  
            UpperCAmelCase__			:     Union[str, Any]										= self.reshape_batch_dim_to_heads(_lowerCAmelCase		)
            UpperCAmelCase__			:     List[str]										= self.proj_attn(_lowerCAmelCase		)
            return self.dropout_layer(_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     (		nn.Module		):
       __lowerCamelCase										=     42
       __lowerCamelCase										=     42
       __lowerCamelCase										=     42
       __lowerCamelCase										=     0.0
       __lowerCamelCase										=     False
       __lowerCamelCase										=     jnp.floataa
       __lowerCamelCase										=     False
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            # self attention (or cross_attention if only_cross_attention is True)
            UpperCAmelCase__			:     List[Any]										= FlaxAttention(
                self.dim  ,	self.n_heads  ,	self.d_head  ,	self.dropout  ,	self.use_memory_efficient_attention  ,	dtype=self.dtype		)
            # cross attention
            UpperCAmelCase__			:     List[str]										= FlaxAttention(
                self.dim  ,	self.n_heads  ,	self.d_head  ,	self.dropout  ,	self.use_memory_efficient_attention  ,	dtype=self.dtype		)
            UpperCAmelCase__			:     Any										= FlaxFeedForward(dim=self.dim  ,	dropout=self.dropout  ,	dtype=self.dtype		)
            UpperCAmelCase__			:     Tuple										= nn.LayerNorm(epsilon=1e-5  ,	dtype=self.dtype		)
            UpperCAmelCase__			:     Optional[Any]										= nn.LayerNorm(epsilon=1e-5  ,	dtype=self.dtype		)
            UpperCAmelCase__			:     Any										= nn.LayerNorm(epsilon=1e-5  ,	dtype=self.dtype		)
            UpperCAmelCase__			:     Union[str, Any]										= nn.Dropout(rate=self.dropout		)
  
  
  
  
       def __call__(							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase=True		):
            # self attention
            UpperCAmelCase__			:     List[str]										= hidden_states
            if self.only_cross_attention:
                 UpperCAmelCase__			:     Dict										= self.attna(self.norma(_lowerCAmelCase		)  ,	_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
            else:
                 UpperCAmelCase__			:     Any										= self.attna(self.norma(_lowerCAmelCase		)  ,	deterministic=_lowerCAmelCase		)
            UpperCAmelCase__			:     List[Any]										= hidden_states + residual
  
            # cross attention
            UpperCAmelCase__			:     Optional[int]										= hidden_states
            UpperCAmelCase__			:     int										= self.attna(self.norma(_lowerCAmelCase		)  ,	_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
            UpperCAmelCase__			:     Any										= hidden_states + residual
  
            # feed forward
            UpperCAmelCase__			:     Union[str, Any]										= hidden_states
            UpperCAmelCase__			:     str										= self.ff(self.norma(_lowerCAmelCase		)  ,	deterministic=_lowerCAmelCase		)
            UpperCAmelCase__			:     int										= hidden_states + residual
  
            return self.dropout_layer(_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     (		nn.Module		):
       __lowerCamelCase										=     42
       __lowerCamelCase										=     42
       __lowerCamelCase										=     42
       __lowerCamelCase										=     1
       __lowerCamelCase										=     0.0
       __lowerCamelCase										=     False
       __lowerCamelCase										=     False
       __lowerCamelCase										=     jnp.floataa
       __lowerCamelCase										=     False
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            UpperCAmelCase__			:     Dict										= nn.GroupNorm(num_groups=32  ,	epsilon=1e-5		)
  
            UpperCAmelCase__			:     Dict										= self.n_heads * self.d_head
            if self.use_linear_projection:
                 UpperCAmelCase__			:     Tuple										= nn.Dense(_lowerCAmelCase  ,	dtype=self.dtype		)
            else:
                 UpperCAmelCase__			:     Dict										= nn.Conv(
                     _lowerCAmelCase  ,	kernel_size=(1, 1)  ,	strides=(1, 1)  ,	padding="""VALID"""  ,	dtype=self.dtype  ,	)
  
            UpperCAmelCase__			:     int										= [
                FlaxBasicTransformerBlock(
                    _lowerCAmelCase  ,	self.n_heads  ,	self.d_head  ,	dropout=self.dropout  ,	only_cross_attention=self.only_cross_attention  ,	dtype=self.dtype  ,	use_memory_efficient_attention=self.use_memory_efficient_attention  ,	)
                for _ in range(self.depth		)
            ]
  
            if self.use_linear_projection:
                 UpperCAmelCase__			:     Union[str, Any]										= nn.Dense(_lowerCAmelCase  ,	dtype=self.dtype		)
            else:
                 UpperCAmelCase__			:     Union[str, Any]										= nn.Conv(
                     _lowerCAmelCase  ,	kernel_size=(1, 1)  ,	strides=(1, 1)  ,	padding="""VALID"""  ,	dtype=self.dtype  ,	)
  
            UpperCAmelCase__			:     Union[str, Any]										= nn.Dropout(rate=self.dropout		)
  
  
  
  
       def __call__(							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase=True		):
            UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			,		UpperCAmelCase__			:     Optional[int]										= hidden_states.shape
            UpperCAmelCase__			:     Union[str, Any]										= hidden_states
            UpperCAmelCase__			:     Union[str, Any]										= self.norm(_lowerCAmelCase		)
            if self.use_linear_projection:
                 UpperCAmelCase__			:     int										= hidden_states.reshape(_lowerCAmelCase  ,	height * width  ,	_lowerCAmelCase		)
                 UpperCAmelCase__			:     int										= self.proj_in(_lowerCAmelCase		)
            else:
                 UpperCAmelCase__			:     Any										= self.proj_in(_lowerCAmelCase		)
                 UpperCAmelCase__			:     List[str]										= hidden_states.reshape(_lowerCAmelCase  ,	height * width  ,	_lowerCAmelCase		)
  
            for transformer_block in self.transformer_blocks:
                 UpperCAmelCase__			:     str										= transformer_block(_lowerCAmelCase  ,	_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
  
            if self.use_linear_projection:
                 UpperCAmelCase__			:     Optional[int]										= self.proj_out(_lowerCAmelCase		)
                 UpperCAmelCase__			:     Tuple										= hidden_states.reshape(_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase		)
            else:
                 UpperCAmelCase__			:     str										= hidden_states.reshape(_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase		)
                 UpperCAmelCase__			:     Union[str, Any]										= self.proj_out(_lowerCAmelCase		)
  
            UpperCAmelCase__			:     List[str]										= hidden_states + residual
            return self.dropout_layer(_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     (		nn.Module		):
       __lowerCamelCase										=     42
       __lowerCamelCase										=     0.0
       __lowerCamelCase										=     jnp.floataa
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            # The second linear layer needs to be called
            # net_2 for now to match the index of the Sequential layer
            UpperCAmelCase__			:     Dict										= FlaxGEGLU(self.dim  ,	self.dropout  ,	self.dtype		)
            UpperCAmelCase__			:     Optional[int]										= nn.Dense(self.dim  ,	dtype=self.dtype		)
  
  
  
  
       def __call__(							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase=True		):
            UpperCAmelCase__			:     Optional[int]										= self.net_a(_lowerCAmelCase  ,	deterministic=_lowerCAmelCase		)
            UpperCAmelCase__			:     List[str]										= self.net_a(_lowerCAmelCase		)
            return hidden_states
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     (		nn.Module		):
       __lowerCamelCase										=     42
       __lowerCamelCase										=     0.0
       __lowerCamelCase										=     jnp.floataa
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            UpperCAmelCase__			:     List[Any]										= self.dim * 4
            UpperCAmelCase__			:     Any										= nn.Dense(inner_dim * 2  ,	dtype=self.dtype		)
            UpperCAmelCase__			:     Dict										= nn.Dropout(rate=self.dropout		)
  
  
  
  
       def __call__(							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase=True		):
            UpperCAmelCase__			:     List[Any]										= self.proj(_lowerCAmelCase		)
            UpperCAmelCase__			,		UpperCAmelCase__			:     int										= jnp.split(_lowerCAmelCase  ,	2  ,	axis=2		)
            return self.dropout_layer(hidden_linear * nn.gelu(_lowerCAmelCase		)  ,	deterministic=_lowerCAmelCase		)
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  
  import collections
  
  from .utils import ExplicitEnum, is_torch_available, logging
  
  
  if is_torch_available():
        import torch
  
  
  SCREAMING_SNAKE_CASE__		:       Any   							=							logging.get_logger(__name__)
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     :
  
  
  
  
       def __init__(							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase=21  ,	_lowerCAmelCase=[]  ,	_lowerCAmelCase=None		):
            UpperCAmelCase__			:     Union[str, Any]										= model
            UpperCAmelCase__			:     Union[str, Any]										= trace_batch_nums
            UpperCAmelCase__			:     Optional[Any]										= abort_after_batch_num
  
            # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
            UpperCAmelCase__			:     int										= collections.deque([]  ,	_lowerCAmelCase		)
            UpperCAmelCase__			:     str										= []
            UpperCAmelCase__			:     Optional[int]										= 0
            UpperCAmelCase__			:     int										= 0
            UpperCAmelCase__			:     int										= False
            UpperCAmelCase__			:     List[str]										= """                 """
  
            self.analyse_model()
  
            self.register_forward_hook()
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase=None		):
            if frame is not None:
                 self.expand_frame(_lowerCAmelCase		)
            self.frames.append("""\n""".join(self.frame		)		)
            UpperCAmelCase__			:     str										= []  # start a new frame
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase		):
            self.frame.append(_lowerCAmelCase		)
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            print("""\n""".join(self.frames		)		)
            UpperCAmelCase__			:     List[Any]										= []
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            UpperCAmelCase__			:     Optional[int]										= []
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            print(f"\nDetected inf/nan during batch_number={self.batch_number}"		)
            print(f"Last {len(self.frames		)} forward frames:"		)
            print(f"{'abs min':8} {'abs max':8} metadata"		)
            print("""\n""".join(self.frames		)		)
            print("""\n\n"""		)
            UpperCAmelCase__			:     Any										= []
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            # extract the fully qualified module names, to be able to report at run time. e.g.:
            # encoder.block.2.layer.0.SelfAttention.o
            #
            # for shared weights only the first shared module name will be registered
            UpperCAmelCase__			:     str										= {m: name for name, m in self.model.named_modules()}
            # self.longest_module_name = max(len(v) for v in self.module_names.values())
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase		):
            if torch.is_tensor(_lowerCAmelCase		):
                 self.expand_frame(get_abs_min_max(_lowerCAmelCase  ,	_lowerCAmelCase		)		)
                 if detect_overflow(_lowerCAmelCase  ,	_lowerCAmelCase		):
                      UpperCAmelCase__			:     Optional[int]										= True
            elif var is None:
                 self.expand_frame(f"{'None':>17} {ctx}"		)
            else:
                 self.expand_frame(f"{'not a tensor':>17} {ctx}"		)
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            self.expand_frame(f"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***"		)
            self.expand_frame(f"{'abs min':8} {'abs max':8} metadata"		)
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"		)
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase		):
            self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}"		)
  
            # params
            for name, p in module.named_parameters(recurse=_lowerCAmelCase		):
                 self.analyse_variable(_lowerCAmelCase  ,	_lowerCAmelCase		)
  
            # inputs
            if isinstance(_lowerCAmelCase  ,	_lowerCAmelCase		):
                 for i, x in enumerate(_lowerCAmelCase		):
                      self.analyse_variable(_lowerCAmelCase  ,	f"input[{i}]"		)
            else:
                 self.analyse_variable(_lowerCAmelCase  ,	"""input"""		)
  
            # outputs
            if isinstance(_lowerCAmelCase  ,	_lowerCAmelCase		):
                 for i, x in enumerate(_lowerCAmelCase		):
                      # possibly a tuple of tuples
                      if isinstance(_lowerCAmelCase  ,	_lowerCAmelCase		):
                           for j, y in enumerate(_lowerCAmelCase		):
                                self.analyse_variable(_lowerCAmelCase  ,	f"output[{i}][{j}]"		)
                      else:
                           self.analyse_variable(_lowerCAmelCase  ,	f"output[{i}]"		)
            else:
                 self.analyse_variable(_lowerCAmelCase  ,	"""output"""		)
  
            self.save_frame()
  
  
  
  
       def      __UpperCAmelCase     (							self		):
            self.model.apply(self._register_forward_hook		)
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase		):
            module.register_forward_hook(self.forward_hook		)
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase		):
            # - input is a tuple of packed inputs (could be non-Tensors)
            # - output could be a Tensor or a tuple of Tensors and non-Tensors
  
            UpperCAmelCase__			:     Union[str, Any]										= False
  
            UpperCAmelCase__			:     List[str]										= True if self.batch_number in self.trace_batch_nums else False
            if trace_mode:
                 self.reset_saved_frames()
  
            if self.total_calls == 0:
                 self.batch_start_frame()
            self.total_calls += 1
  
            # count batch numbers - the very first forward hook of the batch will be called when the
            # batch completes - i.e. it gets called very last - we know this batch has finished
            if module == self.model:
                 self.batch_number += 1
                 UpperCAmelCase__			:     Dict										= True
  
            self.create_frame(_lowerCAmelCase  ,	_lowerCAmelCase  ,	_lowerCAmelCase		)
  
            # if last_frame_of_batch:
            #     self.batch_end_frame()
  
            if trace_mode:
                 self.trace_frames()
  
            if last_frame_of_batch:
                 self.batch_start_frame()
  
            if self.detected_overflow and not trace_mode:
                 self.dump_saved_frames()
  
                 # now we can abort, as it's pointless to continue running
                 raise ValueError(
                     """DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. """
                     """Please scroll up above this traceback to see the activation values prior to this event."""		)
  
            # abort after certain batch if requested to do so
            if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
                 raise ValueError(
                     f"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"
                     f" `abort_after_batch_num={self.abort_after_batch_num}` arg"		)
  
  
  
  
  
  def    _lowerCamelCase	(	__lowerCamelCase   ,       __lowerCamelCase					)   -> Optional[Any]:
  
       '''simple docstring'''
       UpperCAmelCase__			:     Dict										= var.abs()
       return F"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"
  
  
  
  
  
  def    _lowerCamelCase	(	__lowerCamelCase   ,       __lowerCamelCase					)   -> int:
  
       '''simple docstring'''
       UpperCAmelCase__			:     str										= False
       if torch.isnan(__lowerCamelCase					).any().item():
            UpperCAmelCase__			:     List[Any]										= True
            print(F"{ctx} has nans"					)
       if torch.isinf(__lowerCamelCase					).any().item():
            UpperCAmelCase__			:     List[str]										= True
            print(F"{ctx} has infs"					)
  
       # if needed to monitor large elements can enable the following
       if 0:  # and detected:
            UpperCAmelCase__			:     Dict										= var[torch.ge(var.abs()   ,       100					)]
            if naaa.numel() > 0:
                 print(F"{ctx}:  n100={naaa.numel()}"					)
            UpperCAmelCase__			:     Tuple										= var[torch.ge(var.abs()   ,       1000					)]
            if naaaa.numel() > 0:
                 print(F"{ctx}: n1000={naaaa.numel()}"					)
            UpperCAmelCase__			:     Optional[int]										= var[torch.ge(var.abs()   ,       1_0000					)]
            if naaaaa.numel() > 0:
                 print(F"{ctx}: n10000={naaaaa.numel()}"					)
  
       if 0:
            print(F"min={var.min():9.2e} max={var.max():9.2e}"					)
  
       if 0:
            print(F"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})"					)
  
       return detected
  
  
  
  
  
  
  
  class 							UpperCAmelCase_     (		__lowerCamelCase		):
       __lowerCamelCase										=     'underflow_overflow'
       __lowerCamelCase										=     'tpu_metrics_debug'
  
  
  
  
  
  
  