





'''simple docstring'''

import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate

def     __magic_name__							(				__UpperCAmelCase   )    ->       str:

   '''simple docstring'''







   return x + 2







class __a			( unittest.TestCase      ):





   def   UpperCAmelCase__  (       self					:      List[str]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""x = 3"""
      __SCREAMING_SNAKE_CASE   				=							{}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      assert result == 3
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3}     )

      __SCREAMING_SNAKE_CASE   				=							"""x = y"""
      __SCREAMING_SNAKE_CASE   				=							{"""y""": 5}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      # evaluate returns the value of the last assignment.
      assert result == 5
      self.assertDictEqual(lowerCamelCase					,{"""x""": 5, """y""": 5}     )





   def   UpperCAmelCase__  (       self					:      List[str]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""y = add_two(x)"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{"""add_two""": add_two}					,state=lowerCamelCase     )
      assert result == 5
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """y""": 5}     )

      # Won't work without the tool
      with CaptureStdout() as out:
         __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      assert result is None
      assert "tried to execute add_two" in out.out





   def   UpperCAmelCase__  (       self					:      Any     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""x = 3"""
      __SCREAMING_SNAKE_CASE   				=							{}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      assert result == 3
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3}     )





   def   UpperCAmelCase__  (       self					:      Union[str, Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""test_dict = {'x': x, 'y': add_two(x)}"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{"""add_two""": add_two}					,state=lowerCamelCase     )
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """y""": 5}     )
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}     )





   def   UpperCAmelCase__  (       self					:      Dict     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""x = 3\ny = 5"""
      __SCREAMING_SNAKE_CASE   				=							{}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      # evaluate returns the value of the last assignment.
      assert result == 5
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """y""": 5}     )





   def   UpperCAmelCase__  (       self					:      int     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""text = f'This is x: {x}.'"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      # evaluate returns the value of the last assignment.
      assert result == "This is x: 3."
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """text""": """This is x: 3."""}     )





   def   UpperCAmelCase__  (       self					:      Union[str, Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""if x <= 3:\n    y = 2\nelse:\n    y = 5"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      # evaluate returns the value of the last assignment.
      assert result == 2
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """y""": 2}     )

      __SCREAMING_SNAKE_CASE   				=							{"""x""": 8}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      # evaluate returns the value of the last assignment.
      assert result == 5
      self.assertDictEqual(lowerCamelCase					,{"""x""": 8, """y""": 5}     )





   def   UpperCAmelCase__  (       self					:      Dict     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""test_list = [x, add_two(x)]"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{"""add_two""": add_two}					,state=lowerCamelCase     )
      self.assertListEqual(lowerCamelCase					,[3, 5]     )
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """test_list""": [3, 5]}     )





   def   UpperCAmelCase__  (       self					:      List[str]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""y = x"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{}					,state=lowerCamelCase     )
      assert result == 3
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """y""": 3}     )





   def   UpperCAmelCase__  (       self					:      Tuple     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""test_list = [x, add_two(x)]\ntest_list[1]"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{"""add_two""": add_two}					,state=lowerCamelCase     )
      assert result == 5
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """test_list""": [3, 5]}     )

      __SCREAMING_SNAKE_CASE   				=							"""test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"""
      __SCREAMING_SNAKE_CASE   				=							{"""x""": 3}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{"""add_two""": add_two}					,state=lowerCamelCase     )
      assert result == 5
      self.assertDictEqual(lowerCamelCase					,{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}     )





   def   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							"""x = 0\nfor i in range(3):\n    x = i"""
      __SCREAMING_SNAKE_CASE   				=							{}
      __SCREAMING_SNAKE_CASE   				=							evaluate(lowerCamelCase					,{"""range""": range}					,state=lowerCamelCase     )
      assert result == 2
      self.assertDictEqual(lowerCamelCase					,{"""x""": 2, """i""": 2}     )








'''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
UpperCamelCase_    						=       object()

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






def 				_lowerCAmelCase			(   __magic_name__						:     int   , __magic_name__						:     int						)				->  Optional[int]:
						lowercase :							Optional[Any]    		=tuple((re.compile(x + '''$'''						) for x in qs)						)
						for i in range(len(__magic_name__						) - len(__magic_name__						) + 1						):
												lowercase :							Union[str, Any]    		=[x.match(__magic_name__						) for x, y in zip(__magic_name__   , ks[i:]						)]
												if matches and all(__magic_name__						):
																		return True
						return False






def 				_lowerCAmelCase			(   __magic_name__						:     Dict						)				->  List[str]:
						def replace(__magic_name__						:     Optional[int]   , __magic_name__						:     Optional[Any]						):
												for rule, replacement in rules:
																		if _match(__magic_name__   , __magic_name__						):
																								return replacement
												return val

						return replace






def 				_lowerCAmelCase			(   )				->  int:
						return [
						    # embeddings
						    (("transformer", "wpe", "embedding"), P('''mp'''   , __magic_name__						)),
						    (("transformer", "wte", "embedding"), P('''mp'''   , __magic_name__						)),
						    # atention
						    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__magic_name__   , '''mp'''						)),
						    (("attention", "out_proj", "kernel"), P('''mp'''   , __magic_name__						)),
						    (("attention", "out_proj", "bias"), None),
						    # mlp
						    (("mlp", "c_fc", "kernel"), P(__magic_name__   , '''mp'''						)),
						    (("mlp", "c_fc", "bias"), P('''mp'''						)),
						    (("mlp", "c_proj", "kernel"), P('''mp'''   , __magic_name__						)),
						    (("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 				_lowerCAmelCase			(   __magic_name__						:     str						)				->  int:
						lowercase :							int    		=_get_partition_rules()
						lowercase :							Tuple    		=_replacement_rules(__magic_name__						)
						lowercase :							Any    		={k: _unmatched for k in flatten_dict(__magic_name__						)}
						lowercase :							Any    		={k: replace(__magic_name__   , __magic_name__						) for k, v in initd.items()}
						assert _unmatched not in result.values(), "Incomplete partition spec."
						return freeze(unflatten_dict(__magic_name__						)						)





