

"""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       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__   )			->						Union[str, Any]:
 lowerCAmelCase__   :				Optional[Any] 			=				tuple((re.compile(x + "$"   ) for x in qs)   )
 for i in range(len(_lowerCAmelCase   ) - len(_lowerCAmelCase   ) + 1   ):
  lowerCAmelCase__   :				List[Any] 			=				[x.match(_lowerCAmelCase   ) for x, y in zip(_lowerCAmelCase					,					ks[i:]   )]
  if matches and all(_lowerCAmelCase   ):
   return True
 return False



def       SCREAMING_SNAKE_CASE							( lowercase__   )			->						List[str]:
 def replace(lowercase__					,					lowercase__   ):
  for rule, replacement in rules:
   if _match(_lowerCAmelCase					,					_lowerCAmelCase   ):
    return replacement
  return val

 return replace



def       SCREAMING_SNAKE_CASE							( )			->						Optional[Any]:
 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							( lowercase__   )			->						Dict:
 lowerCAmelCase__   :				Optional[int] 			=				_get_partition_rules()
 lowerCAmelCase__   :				Tuple 			=				_replacement_rules(_lowerCAmelCase   )
 lowerCAmelCase__   :				int 			=				{k: _unmatched for k in flatten_dict(_lowerCAmelCase   )}
 lowerCAmelCase__   :				Union[str, Any] 			=				{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

def   __snake_case    (       _lowerCAmelCase			:			list[int]					,					_lowerCAmelCase			:			int						)			->    list[list[int]]:
   A_       :					list[list[int]]        =      []
   A_       :					list[int]        =      []
   A_       :					Dict        =      0
   A_       :					str        =      sum(_lowerCAmelCase						)
   create_state_space_tree(_lowerCAmelCase					,					_lowerCAmelCase					,					_lowerCAmelCase					,					_lowerCAmelCase					,					_lowerCAmelCase					,					_lowerCAmelCase						)
   return result

def   __snake_case    (       _lowerCAmelCase			:			list[int]					,					_lowerCAmelCase			:			int					,					_lowerCAmelCase			:			int					,					_lowerCAmelCase			:			list[int]					,					_lowerCAmelCase			:			list[list[int]]					,					_lowerCAmelCase			:			int					,					)			->    None:
   if sum(_lowerCAmelCase						) > max_sum or (remaining_nums_sum + sum(_lowerCAmelCase						)) < max_sum:
      return
   if sum(_lowerCAmelCase						) == max_sum:
      result.append(_lowerCAmelCase						)
      return
   for index in range(_lowerCAmelCase					,					len(_lowerCAmelCase						)						):
      create_state_space_tree(
          _lowerCAmelCase					,					_lowerCAmelCase					,					index + 1					,					[*path, nums[index]]					,					_lowerCAmelCase					,					remaining_nums_sum - nums[index]					,					)
_lowerCAmelCase				:   List[Any]						       =						[3, 34, 4, 12, 5, 2]
_lowerCAmelCase				:   Dict						       =						9
_lowerCAmelCase				:   Dict						       =						generate_sum_of_subsets_soln(nums, max_sum)





print(*result)
