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

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





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





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):
    def replace(__lowerCAmelCase     ,		__lowerCAmelCase				):
        for rule, replacement in rules:
            if _match(__lowerCAmelCase     ,		__lowerCAmelCase				):
                return replacement
        return val

    return replace





def 						SCREAMING_SNAKE_CASE__     (   ):
    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 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):
    snake_case__    =					_get_partition_rules()
    snake_case__    =					_replacement_rules(__lowerCAmelCase				)
    snake_case__    =					{k: _unmatched for k in flatten_dict(__lowerCAmelCase				)}
    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				)				)





from __future__ import annotations

import math





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):

    if 1 < number < 4:
        # 2 and 3 are primes
        return True
    elif number < 2 or number % 2 == 0 or number % 3 == 0:
        # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
        return False

    # All primes number are in format of 6k +/- 1
    for i in range(5     ,		int(math.sqrt(__lowerCAmelCase				) + 1				)     ,		6				):
        if number % i == 0 or number % (i + 2) == 0:
            return False
    return True





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):
    snake_case__    =					str(__lowerCAmelCase				)
    snake_case__    =					[n]
    for i in range(1     ,		len(__lowerCAmelCase				)				):
        list_nums.append(int(str_num[i:]				)				)
        list_nums.append(int(str_num[:-i]				)				)
    return list_nums





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):
    if len(str(__lowerCAmelCase				)				) > 3:
        if not is_prime(int(str(__lowerCAmelCase				)[-3:]				)				) or not is_prime(int(str(__lowerCAmelCase				)[:3]				)				):
            return False
    return True





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase = 11				):
    snake_case__    =					[]
    snake_case__    =					13
    while len(__lowerCAmelCase				) != count:
        if validate(__lowerCAmelCase				):
            snake_case__    =					list_truncated_nums(__lowerCAmelCase				)
            if all(is_prime(__lowerCAmelCase				) for i in list_nums				):
                list_truncated_primes.append(__lowerCAmelCase				)
        num += 2
    return list_truncated_primes





def 						SCREAMING_SNAKE_CASE__     (   ):
    return sum(compute_truncated_primes(11				)				)


if __name__ == "__main__":
     print(F'''{sum(compute_truncated_primes(11)) = }''')



