
from collections import defaultdict

class       a		:







							def __init__(						self  :  List[str]      ,				snake_case__  :  Optional[Any]      ,				snake_case__  :  Any ):
								"""simple docstring"""







								__lowerCAmelCase         =							total  # total no of tasks (N)

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

								__lowerCAmelCase         =							defaultdict(snake_case__ )  # 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
								__lowerCAmelCase         =							(1 << len(snake_case__ )) - 1







							def       UpperCAmelCase__					(						self  :  Optional[Any]      ,				snake_case__  :  List[str]      ,				snake_case__  :  str ):
								"""simple docstring"""







								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
								__lowerCAmelCase         =							self.count_ways_until(snake_case__      ,				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.
								__lowerCAmelCase         =							total_ways_util

								return self.dp[mask][task_no]







							def       UpperCAmelCase__					(						self  :  Optional[int]      ,				snake_case__  :  Optional[Any] ):
								"""simple docstring"""







								for i in range(len(snake_case__ ) ):
									for j in task_performed[i]:
										self.task[j].append(snake_case__ )

        # 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__":
					UpperCamelCase_   	=	5  # total no of tasks (the value of N)

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








import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
					import jax
					import jax.numpy as jnp

					from transformers.generation import (
					    FlaxForcedBOSTokenLogitsProcessor,
					    FlaxForcedEOSTokenLogitsProcessor,
					    FlaxLogitsProcessorList,
					    FlaxMinLengthLogitsProcessor,
					    FlaxTemperatureLogitsWarper,
					    FlaxTopKLogitsWarper,
					    FlaxTopPLogitsWarper,
					)

@require_flax
class       a		(   unittest.TestCase					):







							def       UpperCAmelCase__					(						self  :  Tuple      ,				snake_case__  :  int      ,				snake_case__  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							jnp.ones((batch_size, length) ) / length
								return scores







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								__lowerCAmelCase         =							20

								__lowerCAmelCase         =							self._get_uniform_logits(batch_size=2      ,				length=snake_case__ )

								# tweak scores to not be uniform anymore
								__lowerCAmelCase         =							scores.at[1, 5].set((1 / length) + 0.1 )  # peak, 1st batch
								__lowerCAmelCase         =							scores.at[1, 10].set((1 / length) - 0.4 )  # valley, 1st batch

								# compute softmax
								__lowerCAmelCase         =							jax.nn.softmax(snake_case__      ,				axis=-1 )

								__lowerCAmelCase         =							FlaxTemperatureLogitsWarper(temperature=0.5 )
								__lowerCAmelCase         =							FlaxTemperatureLogitsWarper(temperature=1.3 )

								__lowerCAmelCase         =							jax.nn.softmax(temp_dist_warper_sharper(snake_case__      ,				scores.copy()      ,				cur_len=snake_case__ )      ,				axis=-1 )
								__lowerCAmelCase         =							jax.nn.softmax(temp_dist_warper_smoother(snake_case__      ,				scores.copy()      ,				cur_len=snake_case__ )      ,				axis=-1 )

								# uniform distribution stays uniform
								self.assertTrue(jnp.allclose(probs[0, :]      ,				warped_prob_sharp[0, :]      ,				atol=1E-3 ) )
								self.assertTrue(jnp.allclose(probs[0, :]      ,				warped_prob_smooth[0, :]      ,				atol=1E-3 ) )

								# sharp peaks get higher, valleys get lower
								self.assertLess(probs[1, :].max()      ,				warped_prob_sharp[1, :].max() )
								self.assertGreater(probs[1, :].min()      ,				warped_prob_sharp[1, :].min() )

								# smooth peaks get lower, valleys get higher
								self.assertGreater(probs[1, :].max()      ,				warped_prob_smooth[1, :].max() )
								self.assertLess(probs[1, :].min()      ,				warped_prob_smooth[1, :].min() )







							def       UpperCAmelCase__					(						self  :  str ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							2

								# create ramp distribution
								__lowerCAmelCase         =							np.broadcast_to(np.arange(snake_case__ )[None, :]      ,				(batch_size, vocab_size) ).copy()
								__lowerCAmelCase         =							ramp_logits[1:, : vocab_size // 2] + vocab_size

								__lowerCAmelCase         =							FlaxTopKLogitsWarper(3 )

								__lowerCAmelCase         =							top_k_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )

								# check that correct tokens are filtered
								self.assertListEqual(jnp.isinf(scores[0] ).tolist()      ,				7 * [True] + 3 * [False] )
								self.assertListEqual(jnp.isinf(scores[1] ).tolist()      ,				2 * [True] + 3 * [False] + 5 * [True] )

								# check special case
								__lowerCAmelCase         =							5
								__lowerCAmelCase         =							FlaxTopKLogitsWarper(top_k=1      ,				filter_value=0.0      ,				min_tokens_to_keep=3 )

								__lowerCAmelCase         =							np.broadcast_to(np.arange(snake_case__ )[None, :]      ,				(batch_size, length) ).copy()
								__lowerCAmelCase         =							top_k_warp_safety_check(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )

								# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
								self.assertListEqual((scores == 0.0).sum(axis=-1 ).tolist()      ,				[2, 2] )







							def       UpperCAmelCase__					(						self  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							2

								# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
								__lowerCAmelCase         =							np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]] ) )

								__lowerCAmelCase         =							FlaxTopPLogitsWarper(0.8 )
								__lowerCAmelCase         =							np.exp(top_p_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ ) )

								# dist should be filtered to keep min num values so that sum is >= top_p
								# exp (-inf) => 0
								__lowerCAmelCase         =							np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]] )
								self.assertTrue(np.allclose(snake_case__      ,				snake_case__      ,				atol=1E-3 ) )

								# check edge cases with negative and extreme logits
								__lowerCAmelCase         =							np.broadcast_to(np.arange(snake_case__ )[None, :]      ,				(batch_size, vocab_size) ).copy() - (
								    vocab_size // 2
								)

								# make ramp_logits more extreme
								__lowerCAmelCase         =							ramp_logits[1] * 1_0_0.0

								# make sure at least 2 tokens are kept
								__lowerCAmelCase         =							FlaxTopPLogitsWarper(0.9      ,				min_tokens_to_keep=2      ,				filter_value=0.0 )
								__lowerCAmelCase         =							top_p_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )

								# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
								self.assertListEqual((filtered_dist != 0.0).sum(axis=-1 ).tolist()      ,				[3, 2] )







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							20
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							0

								__lowerCAmelCase         =							FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=snake_case__ )

								# check that min length is applied at length 5
								__lowerCAmelCase         =							ids_tensor((batch_size, 20)      ,				vocab_size=20 )
								__lowerCAmelCase         =							5
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							min_dist_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()      ,				4 * [-float("inf" )] )

								# check that min length is not applied anymore at length 15
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							15
								__lowerCAmelCase         =							min_dist_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertFalse(jnp.isinf(snake_case__ ).any() )







							def       UpperCAmelCase__					(						self  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							20
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							0

								__lowerCAmelCase         =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )

								# check that all scores are -inf except the bos_token_id score
								__lowerCAmelCase         =							ids_tensor((batch_size, 1)      ,				vocab_size=20 )
								__lowerCAmelCase         =							1
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							logits_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :] ).all() )
								self.assertListEqual(scores[:, bos_token_id].tolist()      ,				4 * [0] )  # score for bos_token_id shold be zero

								# check that bos_token_id is not forced if current length is greater than 1
								__lowerCAmelCase         =							3
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							logits_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertFalse(jnp.isinf(snake_case__ ).any() )







							def       UpperCAmelCase__					(						self  :  str ):
								"""simple docstring"""







								__lowerCAmelCase         =							20
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							0
								__lowerCAmelCase         =							5

								__lowerCAmelCase         =							FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,				eos_token_id=snake_case__ )

								# check that all scores are -inf except the eos_token_id when max_length is reached
								__lowerCAmelCase         =							ids_tensor((batch_size, 4)      ,				vocab_size=20 )
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							logits_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :] ).all() )
								self.assertListEqual(scores[:, eos_token_id].tolist()      ,				4 * [0] )  # score for eos_token_id should be zero

								# check that eos_token_id is not forced if max_length is not reached
								__lowerCAmelCase         =							3
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							logits_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertFalse(jnp.isinf(snake_case__ ).any() )







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							15
								__lowerCAmelCase         =							2
								__lowerCAmelCase         =							1
								__lowerCAmelCase         =							15

								# dummy input_ids and scores
								__lowerCAmelCase         =							ids_tensor((batch_size, sequence_length)      ,				snake_case__ )
								__lowerCAmelCase         =							input_ids.copy()

								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							scores.copy()

								# instantiate all dist processors
								__lowerCAmelCase         =							FlaxTemperatureLogitsWarper(temperature=0.5 )
								__lowerCAmelCase         =							FlaxTopKLogitsWarper(3 )
								__lowerCAmelCase         =							FlaxTopPLogitsWarper(0.8 )

								# instantiate all logits processors
								__lowerCAmelCase         =							FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,				eos_token_id=snake_case__ )

								__lowerCAmelCase         =							10

								# no processor list
								__lowerCAmelCase         =							temp_dist_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							top_k_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							top_p_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							min_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							bos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							eos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )

								# with processor list
								__lowerCAmelCase         =							FlaxLogitsProcessorList(
								    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
								__lowerCAmelCase         =							processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )

								# scores should be equal
								self.assertTrue(jnp.allclose(snake_case__      ,				snake_case__      ,				atol=1E-3 ) )

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()      ,				input_ids_comp.tolist() )







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							15
								__lowerCAmelCase         =							2
								__lowerCAmelCase         =							1
								__lowerCAmelCase         =							15

								# dummy input_ids and scores
								__lowerCAmelCase         =							ids_tensor((batch_size, sequence_length)      ,				snake_case__ )
								__lowerCAmelCase         =							input_ids.copy()

								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							scores.copy()

								# instantiate all dist processors
								__lowerCAmelCase         =							FlaxTemperatureLogitsWarper(temperature=0.5 )
								__lowerCAmelCase         =							FlaxTopKLogitsWarper(3 )
								__lowerCAmelCase         =							FlaxTopPLogitsWarper(0.8 )

								# instantiate all logits processors
								__lowerCAmelCase         =							FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,				eos_token_id=snake_case__ )

								__lowerCAmelCase         =							10

								# no processor list
								def run_no_processor_list(snake_case__  :  int      ,				snake_case__  :  Any      ,				snake_case__  :  int ):
									__lowerCAmelCase         =							temp_dist_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							top_k_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							top_p_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							min_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							bos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							eos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									return scores

								# with processor list
								def run_processor_list(snake_case__  :  Optional[Any]      ,				snake_case__  :  Optional[int]      ,				snake_case__  :  str ):
									__lowerCAmelCase         =							FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
									__lowerCAmelCase         =							processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									return scores

								__lowerCAmelCase         =							jax.jit(snake_case__ )
								__lowerCAmelCase         =							jax.jit(snake_case__ )

								__lowerCAmelCase         =							jitted_run_no_processor_list(snake_case__      ,				snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							jitted_run_processor_list(snake_case__      ,				snake_case__      ,				snake_case__ )

								# scores should be equal
								self.assertTrue(jnp.allclose(snake_case__      ,				snake_case__      ,				atol=1E-3 ) )

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()      ,				input_ids_comp.tolist() )






