


"""simple docstring"""







import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
       import jax.numpy as jnp

if is_tf_available():
       import tensorflow as tf

if is_torch_available():
       import torch






class       __A				(       unittest.TestCase			):






 """simple docstring"""




 def       SCREAMING_SNAKE_CASE							(		self						)					->     Tuple:
     a 		={
         '''task_specific_params''': {
             '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 128, '''min_length''': 12, '''num_beams''': 4},
             '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 142, '''min_length''': 56, '''num_beams''': 4},
             '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 62, '''min_length''': 11, '''num_beams''': 6},
         }
     }
     a 		={
         '''task_specific_params.summarization.length_penalty''': 1.0,
         '''task_specific_params.summarization.max_length''': 128,
         '''task_specific_params.summarization.min_length''': 12,
         '''task_specific_params.summarization.num_beams''': 4,
         '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
         '''task_specific_params.summarization_cnn.max_length''': 142,
         '''task_specific_params.summarization_cnn.min_length''': 56,
         '''task_specific_params.summarization_cnn.num_beams''': 4,
         '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
         '''task_specific_params.summarization_xsum.max_length''': 62,
         '''task_specific_params.summarization_xsum.min_length''': 11,
         '''task_specific_params.summarization_xsum.num_beams''': 6,
     }

     self.assertEqual(flatten_dict(__A						) ,    __A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     str:
     a 		=np.random.randn(3 ,    4						)
     self.assertTrue(np.allclose(transpose(__A						) ,    x.transpose()						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     self.assertTrue(np.allclose(transpose(__A ,    axes=(1, 2, 0)						) ,    x.transpose((1, 2, 0)						)						)						)


 @require_torch
 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=np.random.randn(3 ,    4						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(transpose(__A						) ,    transpose(__A						).numpy()						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(transpose(__A ,    axes=(1, 2, 0)						) ,    transpose(__A ,    axes=(1, 2, 0)						).numpy()						)						)


 @require_tf
 def       SCREAMING_SNAKE_CASE							(		self						)					->     str:
     a 		=np.random.randn(3 ,    4						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(transpose(__A						) ,    transpose(__A						).numpy()						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(transpose(__A ,    axes=(1, 2, 0)						) ,    transpose(__A ,    axes=(1, 2, 0)						).numpy()						)						)


 @require_flax
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Any:
     a 		=np.random.randn(3 ,    4						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(transpose(__A						) ,    np.asarray(transpose(__A						)						)						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(transpose(__A ,    axes=(1, 2, 0)						) ,    np.asarray(transpose(__A ,    axes=(1, 2, 0)						)						)						)						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Dict:
     a 		=np.random.randn(3 ,    4						)
     self.assertTrue(np.allclose(reshape(__A ,    (4, 3)						) ,    np.reshape(__A ,    (4, 3)						)						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     self.assertTrue(np.allclose(reshape(__A ,    (12, 5)						) ,    np.reshape(__A ,    (12, 5)						)						)						)


 @require_torch
 def       SCREAMING_SNAKE_CASE							(		self						)					->     int:
     a 		=np.random.randn(3 ,    4						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(reshape(__A ,    (4, 3)						) ,    reshape(__A ,    (4, 3)						).numpy()						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(reshape(__A ,    (12, 5)						) ,    reshape(__A ,    (12, 5)						).numpy()						)						)


 @require_tf
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Optional[Any]:
     a 		=np.random.randn(3 ,    4						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(reshape(__A ,    (4, 3)						) ,    reshape(__A ,    (4, 3)						).numpy()						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(reshape(__A ,    (12, 5)						) ,    reshape(__A ,    (12, 5)						).numpy()						)						)


 @require_flax
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Optional[int]:
     a 		=np.random.randn(3 ,    4						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(reshape(__A ,    (4, 3)						) ,    np.asarray(reshape(__A ,    (4, 3)						)						)						)						)

     a 		=np.random.randn(3 ,    4 ,    5						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(reshape(__A ,    (12, 5)						) ,    np.asarray(reshape(__A ,    (12, 5)						)						)						)						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     str:
     a 		=np.random.randn(1 ,    3 ,    4						)
     self.assertTrue(np.allclose(squeeze(__A						) ,    np.squeeze(__A						)						)						)

     a 		=np.random.randn(1 ,    4 ,    1 ,    5						)
     self.assertTrue(np.allclose(squeeze(__A ,    axis=2						) ,    np.squeeze(__A ,    axis=2						)						)						)


 @require_torch
 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=np.random.randn(1 ,    3 ,    4						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(squeeze(__A						) ,    squeeze(__A						).numpy()						)						)

     a 		=np.random.randn(1 ,    4 ,    1 ,    5						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(squeeze(__A ,    axis=2						) ,    squeeze(__A ,    axis=2						).numpy()						)						)


 @require_tf
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Tuple:
     a 		=np.random.randn(1 ,    3 ,    4						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(squeeze(__A						) ,    squeeze(__A						).numpy()						)						)

     a 		=np.random.randn(1 ,    4 ,    1 ,    5						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(squeeze(__A ,    axis=2						) ,    squeeze(__A ,    axis=2						).numpy()						)						)


 @require_flax
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=np.random.randn(1 ,    3 ,    4						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(squeeze(__A						) ,    np.asarray(squeeze(__A						)						)						)						)

     a 		=np.random.randn(1 ,    4 ,    1 ,    5						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(squeeze(__A ,    axis=2						) ,    np.asarray(squeeze(__A ,    axis=2						)						)						)						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=np.random.randn(3 ,    4						)
     self.assertTrue(np.allclose(expand_dims(__A ,    axis=1						) ,    np.expand_dims(__A ,    axis=1						)						)						)


 @require_torch
 def       SCREAMING_SNAKE_CASE							(		self						)					->     str:
     a 		=np.random.randn(3 ,    4						)
     a 		=torch.tensor(__A						)
     self.assertTrue(np.allclose(expand_dims(__A ,    axis=1						) ,    expand_dims(__A ,    axis=1						).numpy()						)						)


 @require_tf
 def       SCREAMING_SNAKE_CASE							(		self						)					->     int:
     a 		=np.random.randn(3 ,    4						)
     a 		=tf.constant(__A						)
     self.assertTrue(np.allclose(expand_dims(__A ,    axis=1						) ,    expand_dims(__A ,    axis=1						).numpy()						)						)





 @require_flax
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=np.random.randn(3 ,    4						)
     a 		=jnp.array(__A						)
     self.assertTrue(np.allclose(expand_dims(__A ,    axis=1						) ,    np.asarray(expand_dims(__A ,    axis=1						)						)						)						)



"""simple docstring"""







import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs




@pytest.mark.parametrize(
    "kwargs, expected"  ,						[
        ({"num_shards": 0, "max_num_jobs": 1}, []),
        ({"num_shards": 10, "max_num_jobs": 1}, [range(10   )]),
        ({"num_shards": 10, "max_num_jobs": 10}, [range(lowerCAmelCase_  ,						i + 1   ) for i in range(10   )]),
        ({"num_shards": 1, "max_num_jobs": 10}, [range(1   )]),
        ({"num_shards": 10, "max_num_jobs": 3}, [range(0  ,						4   ), range(4  ,						7   ), range(7  ,						10   )]),
        ({"num_shards": 3, "max_num_jobs": 10}, [range(0  ,						1   ), range(1  ,						2   ), range(2  ,						3   )]),
    ]  ,						)
def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				_distribute_shards(**lowerCAmelCase_   )
					assert out == expected




@pytest.mark.parametrize(
    "gen_kwargs, max_num_jobs, expected"  ,						[
        ({"foo": 0}, 10, [{"foo": 0}]),
        ({"shards": [0, 1, 2, 3]}, 1, [{"shards": [0, 1, 2, 3]}]),
        ({"shards": [0, 1, 2, 3]}, 4, [{"shards": [0]}, {"shards": [1]}, {"shards": [2]}, {"shards": [3]}]),
        ({"shards": [0, 1]}, 4, [{"shards": [0]}, {"shards": [1]}]),
        ({"shards": [0, 1, 2, 3]}, 2, [{"shards": [0, 1]}, {"shards": [2, 3]}]),
    ]  ,						)
def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				_split_gen_kwargs(lowerCAmelCase_  ,						lowerCAmelCase_   )
					assert out == expected




@pytest.mark.parametrize(
    "gen_kwargs, expected"  ,						[
        ({"foo": 0}, 1),
        ({"shards": [0]}, 1),
        ({"shards": [0, 1, 2, 3]}, 4),
        ({"shards": [0, 1, 2, 3], "foo": 0}, 4),
        ({"shards": [0, 1, 2, 3], "other": (0, 1)}, 4),
        ({"shards": [0, 1, 2, 3], "shards2": [0, 1]}, RuntimeError),
    ]  ,						)
def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_   ):




					'''simple docstring'''


					if expected is RuntimeError:
										with pytest.raises(lowerCAmelCase_   ):
															_number_of_shards_in_gen_kwargs(lowerCAmelCase_   )
					else:
										__SCREAMING_SNAKE_CASE						   =				_number_of_shards_in_gen_kwargs(lowerCAmelCase_   )
										assert out == expected
