





class 					__snake_case	:


    '''simple docstring'''







    def __init__(       self   ,				A_   ,				A_=None   ,				A_=None     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			data
        SCREAMING_SNAKE_CASE__											=			previous
        SCREAMING_SNAKE_CASE__											=			next_node



    def __str__(       self     ):





        '''simple docstring'''

        return f'''{self.data}'''



    def  lowercase_					(       self     ):





        '''simple docstring'''

        return self.data



    def  lowercase_					(       self     ):





        '''simple docstring'''

        return self.next




    def  lowercase_					(       self     ):





        '''simple docstring'''

        return self.previous





class 					__snake_case	:


    '''simple docstring'''







    def __init__(       self   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			head



    def __iter__(       self     ):





        '''simple docstring'''

        return self




    def  lowercase_					(       self     ):





        '''simple docstring'''

        if not self.current:
            raise StopIteration
        else:
            SCREAMING_SNAKE_CASE__											=			self.current.get_data()
            SCREAMING_SNAKE_CASE__											=			self.current.get_next()
            return value





class 					__snake_case	:


    '''simple docstring'''







    def __init__(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			None  # First node in list
        SCREAMING_SNAKE_CASE__											=			None  # Last node in list



    def __str__(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.head
        SCREAMING_SNAKE_CASE__											=			[]
        while current is not None:
            nodes.append(current.get_data()     )
            SCREAMING_SNAKE_CASE__											=			current.get_next()
        return " ".join(str(A_     ) for node in nodes     )



    def __contains__(       self   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.head
        while current:
            if current.get_data() == value:
                return True
            SCREAMING_SNAKE_CASE__											=			current.get_next()
        return False



    def __iter__(       self     ):





        '''simple docstring'''

        return LinkedListIterator(self.head     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        if self.head:
            return self.head.get_data()
        return None



    def  lowercase_					(       self     ):





        '''simple docstring'''

        if self.tail:
            return self.tail.get_data()
        return None



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        if self.head is None:
            SCREAMING_SNAKE_CASE__											=			node
            SCREAMING_SNAKE_CASE__											=			node
        else:
            self.insert_before_node(self.head   ,				A_     )



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        if self.head is None:
            self.set_head(A_     )
        else:
            self.insert_after_node(self.tail   ,				A_     )



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			Node(A_     )
        if self.head is None:
            self.set_head(A_     )
        else:
            self.set_tail(A_     )



    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			node
        SCREAMING_SNAKE_CASE__											=			node.previous

        if node.get_previous() is None:
            SCREAMING_SNAKE_CASE__											=			node_to_insert
        else:
            SCREAMING_SNAKE_CASE__											=			node_to_insert

        SCREAMING_SNAKE_CASE__											=			node_to_insert



    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			node
        SCREAMING_SNAKE_CASE__											=			node.next

        if node.get_next() is None:
            SCREAMING_SNAKE_CASE__											=			node_to_insert
        else:
            SCREAMING_SNAKE_CASE__											=			node_to_insert

        SCREAMING_SNAKE_CASE__											=			node_to_insert



    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			1
        SCREAMING_SNAKE_CASE__											=			Node(A_     )
        SCREAMING_SNAKE_CASE__											=			self.head
        while node:
            if current_position == position:
                self.insert_before_node(A_   ,				A_     )
                return
            current_position += 1
            SCREAMING_SNAKE_CASE__											=			node.next
        self.insert_after_node(self.tail   ,				A_     )



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.head
        while node:
            if node.get_data() == item:
                return node
            SCREAMING_SNAKE_CASE__											=			node.get_next()
        raise Exception('''Node not found'''     )



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        if (node := self.get_node(A_     )) is not None:
            if node == self.head:
                SCREAMING_SNAKE_CASE__											=			self.head.get_next()

            if node == self.tail:
                SCREAMING_SNAKE_CASE__											=			self.tail.get_previous()

            self.remove_node_pointers(A_     )



    @staticmethod
    def  lowercase_					(       A_     ):





        '''simple docstring'''

        if node.get_next():
            SCREAMING_SNAKE_CASE__											=			node.previous

        if node.get_previous():
            SCREAMING_SNAKE_CASE__											=			node.next

        SCREAMING_SNAKE_CASE__											=			None
        SCREAMING_SNAKE_CASE__											=			None




    def  lowercase_					(       self     ):





        '''simple docstring'''

        return self.head is None





def 							__snake_case							(    )				->						None:
    pass


if __name__ == "__main__":
  import doctest

  doctest.testmod()







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
_A					:				Optional[int]				=			object()

# For specifying empty leaf dict `{}`
_A					:				Tuple				=			object()


def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_	)				->						Dict:
    SCREAMING_SNAKE_CASE__											=			tuple((re.compile(x + '''$'''	) for x in qs)	)
    for i in range(len(lowerCAmelCase_	) - len(lowerCAmelCase_	) + 1	):
        SCREAMING_SNAKE_CASE__											=			[x.match(lowerCAmelCase_	) for x, y in zip(lowerCAmelCase_ ,							ks[i:]	)]
        if matches and all(lowerCAmelCase_	):
            return True
    return False


def 							__snake_case							(    lowerCAmelCase_	)				->						int:
    def replace(lowerCAmelCase_ ,							lowerCAmelCase_	):
        for rule, replacement in rules:
            if _match(lowerCAmelCase_ ,							lowerCAmelCase_	):
                return replacement
        return val

    return replace


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