



'''simple docstring'''


import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


_lowerCamelCase	:       List[Any]      	=	logging.get_logger(__name__)







def 				__lowerCamelCase				(   A__    )       ->	List[Any]:






   """simple docstring"""
   UpperCamelCase					=  R'\w+[.]\d+'
   UpperCamelCase					=  re.findall(__snake_case    ,						__snake_case    )
   for pat in pats:
      UpperCamelCase					=  key.replace(__snake_case    ,						'_'.join(pat.split('.'    )    )    )
   return key







def 				__lowerCamelCase				(   A__    ,						A__    ,						A__    )       ->	List[Any]:






   """simple docstring"""
   UpperCamelCase					=  pt_tuple_key[:-1] + ('scale',)
   if (
       any('norm' in str_ for str_ in pt_tuple_key    )
       and (pt_tuple_key[-1] == "bias")
       and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
       and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
   ):
      UpperCamelCase					=  pt_tuple_key[:-1] + ('scale',)
      return renamed_pt_tuple_key, pt_tensor
   elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
      UpperCamelCase					=  pt_tuple_key[:-1] + ('scale',)
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
      UpperCamelCase					=  pt_tuple_key[:-1] + ('embedding',)
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   UpperCamelCase					=  pt_tuple_key[:-1] + ('kernel',)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
      UpperCamelCase					=  pt_tensor.transpose(2    ,						3    ,						1    ,						0    )
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   UpperCamelCase					=  pt_tuple_key[:-1] + ('kernel',)
   if pt_tuple_key[-1] == "weight":
      UpperCamelCase					=  pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   UpperCamelCase					=  pt_tuple_key[:-1] + ('weight',)
   if pt_tuple_key[-1] == "gamma":
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm bias
   UpperCamelCase					=  pt_tuple_key[:-1] + ('bias',)
   if pt_tuple_key[-1] == "beta":
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor







def 				__lowerCamelCase				(   A__    ,						A__    ,						A__=42    )       ->	List[str]:






   """simple docstring"""
   UpperCamelCase					=  {k: v.numpy() for k, v in pt_state_dict.items()}

   # Step 2: Since the model is stateless, get random Flax params
   UpperCamelCase					=  flax_model.init_weights(PRNGKey(__snake_case    )    )

   UpperCamelCase					=  flatten_dict(__snake_case    )
   UpperCamelCase					=  {}

   # Need to change some parameters name to match Flax names
   for pt_key, pt_tensor in pt_state_dict.items():
      UpperCamelCase					=  rename_key(__snake_case    )
      UpperCamelCase					=  tuple(renamed_pt_key.split('.'    )    )

      # Correctly rename weight parameters
      UpperCamelCase				,      UpperCamelCase					=  rename_key_and_reshape_tensor(__snake_case    ,						__snake_case    ,						__snake_case    )

      if flax_key in random_flax_state_dict:
         if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
            raise ValueError(
                F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""    )

        # also add unexpected weight so that warning is thrown
      UpperCamelCase					=  jnp.asarray(__snake_case    )

   return unflatten_dict(__snake_case    )



"""simple docstring"""





from collections import defaultdict






class     _UpperCAmelCase:



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






							'''simple docstring'''
							_UpperCamelCase							 =  total  # total no of tasks (N)

							# DP table will have a dimension of (2^M)*N
							# initially all values are set to -1
							_UpperCamelCase							 =  [
							    [-1 for i in range(total + 1)] for j in range(2 ** len(__a))
							]

							_UpperCamelCase							 =  defaultdict(__a)  # stores the list of persons for each task

							# final_mask is used to check if all persons are included by setting all bits
							# to 1
							_UpperCamelCase							 =  (1 << len(__a)) - 1



		def 						UpperCAmelCase							(					self				,		__a				,		__a)       ->       Dict:






							'''simple docstring'''
							# if mask == self.finalmask all persons are distributed tasks, return 1
							if mask == self.final_mask:
												return 1

							# if not everyone gets the task and no more tasks are available, return 0
							if task_no > self.total_tasks:
												return 0

							# if case already considered
							if self.dp[mask][task_no] != -1:
												return self.dp[mask][task_no]

							# Number of ways when we don't this task in the arrangement
							_UpperCamelCase							 =  self.count_ways_until(__a				,		task_no + 1)

							# now assign the tasks one by one to all possible persons and recursively
							# assign for the remaining tasks.
							if task_no in self.task:
												for p in self.task[task_no]:
																	# if p is already given a task
																	if mask & (1 << p):
																						continue

																	# assign this task to p and change the mask value. And recursively
																	# assign tasks with the new mask value.
																	total_ways_util += self.count_ways_until(mask | (1 << p)				,		task_no + 1)

        # save the value.
							_UpperCamelCase							 =  total_ways_util

							return self.dp[mask][task_no]




		def 						UpperCAmelCase							(					self				,		__a)       ->       int:






							'''simple docstring'''
							# Store the list of persons for each task
							for i in range(len(__a)):
												for j in task_performed[i]:
																	self.task[j].append(__a)

        # call the function to fill the DP table, final answer is stored in dp[0][1]
							return self.count_ways_until(0				,		1)


if __name__ == "__main__":
				_a														=		5  # total no of tasks (the value of N)

				# the list of tasks that can be done by M persons.
				_a														=		[[1, 3, 4], [1, 2, 5], [3, 4]]
				print(
				    AssignmentUsingBitmask(task_performed, total_tasks).count_no_of_ways(
				        task_performed
				    )
				)






