




"""simple docstring"""



import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
      import torch

      from transformers.activations import gelu_new, gelu_python, get_activation




@require_torch
class      UpperCAmelCase_  (						unittest.TestCase							):




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




     '''simple docstring'''

     A__	=    torch.tensor([-100, -1, -0.1, 0, 0.1, 1.0, 100]				)
     A__	=    get_activation("gelu"				)
     self.assertTrue(torch.allclose(gelu_python(snake_case_				)     ,			torch_builtin(snake_case_				)				)				)
     self.assertFalse(torch.allclose(gelu_python(snake_case_				)     ,			gelu_new(snake_case_				)				)				)




    def 							__magic_name__							(   self						:	str				)				->							Dict:




     '''simple docstring'''

     A__	=    torch.tensor([-100, -1, -0.1, 0, 0.1, 1.0, 100]				)
     A__	=    get_activation("gelu"				)
     A__	=    get_activation("gelu_10"				)

     A__	=    torch_builtin(snake_case_				)
     A__	=    geluaa(snake_case_				)

     A__	=    torch.where(y_gelu_aa < 10.0     ,			1     ,			0				)

     self.assertTrue(torch.max(snake_case_				).item() == 10.0				)
     self.assertTrue(torch.allclose(y_gelu * clipped_mask     ,			y_gelu_aa * clipped_mask				)				)




    def 							__magic_name__							(   self						:	List[str]				)				->							str:




     '''simple docstring'''

     get_activation("gelu"				)
     get_activation("gelu_10"				)
     get_activation("gelu_fast"				)
     get_activation("gelu_new"				)
     get_activation("gelu_python"				)
     get_activation("gelu_pytorch_tanh"				)
     get_activation("linear"				)
     get_activation("mish"				)
     get_activation("quick_gelu"				)
     get_activation("relu"				)
     get_activation("sigmoid"				)
     get_activation("silu"				)
     get_activation("swish"				)
     get_activation("tanh"				)
     with self.assertRaises(snake_case_				):
      get_activation("bogus"				)
     with self.assertRaises(snake_case_				):
      get_activation(snake_case_				)




    def 							__magic_name__							(   self						:	int				)				->							Tuple:




     '''simple docstring'''

     A__	=    get_activation("gelu"				)
     A__	=    1
     A__	=    get_activation("gelu"				)
     self.assertEqual(acta.a     ,			1				)
     with self.assertRaises(snake_case_				):
      A__	=    acta.a











"""simple docstring"""



import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
SCREAMING_SNAKE_CASE    =	object()

# For specifying empty leaf dict `{}`
SCREAMING_SNAKE_CASE    =	object()







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		List[Any]:
 A__	=    tuple((re.compile(x + "$"						) for x in qs)						)
 for i in range(len(lowercase_						) - len(lowercase_						) + 1						):
  A__	=    [x.match(lowercase_						) for x, y in zip(lowercase_      ,			ks[i:]						)]
  if matches and all(lowercase_						):
   return True
 return False







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		Optional[int]:
 def replace(lowercase_      ,			lowercase_						):
  for rule, replacement in rules:
   if _match(lowercase_      ,			lowercase_						):
    return replacement
  return val

 return replace







def   _SCREAMING_SNAKE_CASE       (							)						->		List[str]:
 return [
     # embeddings
     (("transformer", "wpe", "embedding"), P("mp"      ,			lowercase_						)),
     (("transformer", "wte", "embedding"), P("mp"      ,			lowercase_						)),
     # atention
     (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(lowercase_      ,			"mp"						)),
     (("attention", "out_proj", "kernel"), P("mp"      ,			lowercase_						)),
     (("attention", "out_proj", "bias"), None),
     # mlp
     (("mlp", "c_fc", "kernel"), P(lowercase_      ,			"mp"						)),
     (("mlp", "c_fc", "bias"), P("mp"						)),
     (("mlp", "c_proj", "kernel"), P("mp"      ,			lowercase_						)),
     (("mlp", "c_proj", "bias"), None),
     # layer norms
     ((r"ln_\d+", "bias"), None),
     ((r"\d+", r"ln_\d+", "scale"), None),
     (("ln_f", "bias"), None),
     (("ln_f", "scale"), None),
 ]







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 A__	=    _get_partition_rules()
 A__	=    _replacement_rules(lowercase_						)
 A__	=    {k: _unmatched for k in flatten_dict(lowercase_						)}
 A__	=    {k: replace(lowercase_      ,			lowercase_						) for k, v in initd.items()}
 assert _unmatched not in result.values(), "Incomplete partition spec."
 return freeze(unflatten_dict(lowercase_						)						)





