





import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch



def 		A_				(     _lowerCAmelCase     )						-> int:  # picklable for multiprocessing
 return x.sum()



def 		A_				(     _lowerCAmelCase     )						-> str:  # picklable for multiprocessing
 return i + 1



@dataclass
class 					A__   :
      _UpperCAmelCase	:int
      _UpperCAmelCase	:str



class 					A__   (       __snake_case			):





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   {}
       UpperCamelCase	:					Any       	=   []
       UpperCamelCase	:					Dict       	=   1
       UpperCamelCase	:					Optional[int]       	=   [1, 2]
       UpperCamelCase	:					Union[str, Any]       	=   {"a": 1, "b": 2}
       UpperCamelCase	:					Optional[Any]       	=   {"a": [1, 2], "b": [3, 4]}
       UpperCamelCase	:					Optional[Any]       	=   {"a": {"1": 1}, "b": 2}
       UpperCamelCase	:					Optional[int]       	=   {"a": 1, "b": 2, "c": 3, "d": 4}
       UpperCamelCase	:					Dict       	=   {}
       UpperCamelCase	:					List[str]       	=   []
       UpperCamelCase	:					Union[str, Any]       	=   2
       UpperCamelCase	:					str       	=   [2, 3]
       UpperCamelCase	:					str       	=   {"a": 2, "b": 3}
       UpperCamelCase	:					Optional[Any]       	=   {"a": [2, 3], "b": [4, 5]}
       UpperCamelCase	:					List[str]       	=   {"a": {"1": 2}, "b": 3}
       UpperCamelCase	:					Dict       	=   {"a": 2, "b": 3, "c": 4, "d": 5}
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_   )       ,							A_   )

       UpperCamelCase	:					Any       	=   2
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(map_nested(A_       ,							A_       ,							num_proc=A_   )       ,							A_   )

       UpperCamelCase	:					Optional[int]       	=   {"a": np.eye(2   ), "b": np.zeros(3   ), "c": np.ones(2   )}
       UpperCamelCase	:					Optional[Any]       	=   {"a": 2, "b": 0, "c": 2}
       UpperCamelCase	:					Optional[Any]       	=   {
           "a": np.eye(2   ).astype(A_   ),
           "b": np.zeros(3   ).astype(A_   ),
           "c": np.ones(2   ).astype(A_   ),
       }
       self.assertEqual(map_nested(A_       ,							A_       ,							map_numpy=A_   )       ,							A_   )
       self.assertEqual(
           {k: v.tolist() for k, v in map_nested(A_       ,							A_       ,							map_numpy=A_   ).items()}       ,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,							)
       self.assertEqual(map_nested(A_       ,							A_       ,							map_numpy=A_       ,							num_proc=A_   )       ,							A_   )
       self.assertEqual(
           {k: v.tolist() for k, v in map_nested(A_       ,							A_       ,							map_numpy=A_       ,							num_proc=A_   ).items()}       ,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,							)
       with self.assertRaises(A_   ):  # can't pickle a local lambda
        map_nested(lambda A_   : x + 1       ,							A_       ,							num_proc=A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   {"a": 1, "b": 2}
       UpperCamelCase	:					Union[str, Any]       	=   {"a": 3, "b": 4}
       UpperCamelCase	:					Optional[int]       	=   {"a": 5, "b": 6}
       UpperCamelCase	:					Tuple       	=   sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]   )
       self.assertEqual(sorted(zip_dict(A_       ,							A_       ,							A_   )   )       ,							A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       class 					A__   :
             _UpperCAmelCase	:int						    =   'bar'


       UpperCamelCase	:					Dict       	=   Foo()
       self.assertEqual(foo.my_attr       ,							"bar"   )
       with temporary_assignment(A_       ,							"my_attr"       ,							"BAR"   ):
        self.assertEqual(foo.my_attr       ,							"BAR"   )
       self.assertEqual(foo.my_attr       ,							"bar"   )



@pytest.mark.parametrize(
    "iterable_length, num_proc, expected_num_proc"							,   [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]							,   )
def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase     )						-> List[Any]:
 with patch("datasets.utils.py_utils._single_map_nested"     ) as mock_single_map_nested, patch(
     "datasets.parallel.parallel.Pool"     ) as mock_multiprocessing_pool:
  UpperCamelCase	:					int       	=   {F"""{i}""": i for i in range(_lowerCAmelCase     )}
  UpperCamelCase	:					Optional[int]       	=   map_nested(lambda _lowerCAmelCase     : x + 10							,   _lowerCAmelCase							,   num_proc=_lowerCAmelCase							,   parallel_min_length=16     )
  if expected_num_proc == 1:
   assert mock_single_map_nested.called
   assert not mock_multiprocessing_pool.called
  else:
   assert not mock_single_map_nested.called
   assert mock_multiprocessing_pool.called
   assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc



class 					A__   (       __snake_case			):





      @require_tf
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       import tensorflow as tf
       from tensorflow.keras import layers

       UpperCamelCase	:					Dict       	=   layers.Dense(2   )

       def gen_random_output():
        UpperCamelCase	:					Optional[int]       	=   tf.random.uniform((1, 3)   )
        return model(A_   ).numpy()

       with temp_seed(42       ,							set_tensorflow=A_   ):
        UpperCamelCase	:					Optional[Any]       	=   gen_random_output()
       with temp_seed(42       ,							set_tensorflow=A_   ):
        UpperCamelCase	:					List[str]       	=   gen_random_output()
       UpperCamelCase	:					Optional[int]       	=   gen_random_output()

       np.testing.assert_equal(A_       ,							A_   )
       self.assertGreater(np.abs(outa - outa   ).sum()       ,							0   )





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       import torch

       def gen_random_output():
        UpperCamelCase	:					Any       	=   torch.nn.Linear(3       ,							2   )
        UpperCamelCase	:					Union[str, Any]       	=   torch.rand(1       ,							3   )
        return model(A_   ).detach().numpy()

       with temp_seed(42       ,							set_pytorch=A_   ):
        UpperCamelCase	:					List[str]       	=   gen_random_output()
       with temp_seed(42       ,							set_pytorch=A_   ):
        UpperCamelCase	:					Dict       	=   gen_random_output()
       UpperCamelCase	:					str       	=   gen_random_output()

       np.testing.assert_equal(A_       ,							A_   )
       self.assertGreater(np.abs(outa - outa   ).sum()       ,							0   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       def gen_random_output():
        return np.random.rand(1       ,							3   )

       with temp_seed(42   ):
        UpperCamelCase	:					Tuple       	=   gen_random_output()
       with temp_seed(42   ):
        UpperCamelCase	:					int       	=   gen_random_output()
       UpperCamelCase	:					Any       	=   gen_random_output()

       np.testing.assert_equal(A_       ,							A_   )
       self.assertGreater(np.abs(outa - outa   ).sum()       ,							0   )



@pytest.mark.parametrize("input_data"							,   [{}]     )
def 		A_				(     _lowerCAmelCase     )						-> Dict:
 UpperCamelCase	:					Union[str, Any]       	=   NestedDataStructure(_lowerCAmelCase     ).data
 assert output_data == input_data



@pytest.mark.parametrize(
    "data, expected_output"							,   [
        ({}, []),
        ([], []),
        ("foo", ["foo"]),
        (["foo", "bar"], ["foo", "bar"]),
        ([["foo", "bar"]], ["foo", "bar"]),
        ([[["foo"], ["bar"]]], ["foo", "bar"]),
        ([[["foo"], "bar"]], ["foo", "bar"]),
        ({"a": 1, "b": 2}, [1, 2]),
        ({"a": [1, 2], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, [4]]}, [1, 2, 3, 4]),
        ({"a": {"1": 1}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": [2]}, [1, 2]),
    ]							,   )
def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase     )						-> Tuple:
 UpperCamelCase	:					Union[str, Any]       	=   NestedDataStructure(_lowerCAmelCase     ).flatten()
 assert output == expected_output



def 		A_				(     )						-> List[Any]:
 UpperCamelCase	:					Dict       	=   A(x=1							,   y="foobar"     )
 UpperCamelCase	:					Optional[Any]       	=   {"x": 1, "y": "foobar"}
 assert asdict(_lowerCAmelCase     ) == expected_output

 UpperCamelCase	:					Tuple       	=   {"a": {"b": A(x=10							,   y="foo"     )}, "c": [A(x=20							,   y="bar"     )]}
 UpperCamelCase	:					str       	=   {"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
 assert asdict(_lowerCAmelCase     ) == expected_output

 with pytest.raises(_lowerCAmelCase     ):
  asdict([1, A(x=10							,   y="foo"     )]     )



def 		A_				(     _lowerCAmelCase     )						-> Any:
 return text.split()



def 		A_				(     _lowerCAmelCase     )						-> Optional[int]:
 yield (time.time(), content)
 time.sleep(2     )
 yield (time.time(), content)



def 		A_				(     )						-> int:
 with Pool(2     ) as pool:
  UpperCamelCase	:					Optional[Any]       	=   list(iflatmap_unordered(_lowerCAmelCase							,   _split_text							,   kwargs_iterable=[{"text": "hello there"}] * 10     )     )
  assert out.count("hello"     ) == 10
  assert out.count("there"     ) == 10
  assert len(_lowerCAmelCase     ) == 20

 # check multiprocess from pathos (uses dill for pickling)
 with multiprocess.Pool(2     ) as pool:
  UpperCamelCase	:					Tuple       	=   list(iflatmap_unordered(_lowerCAmelCase							,   _split_text							,   kwargs_iterable=[{"text": "hello there"}] * 10     )     )
  assert out.count("hello"     ) == 10
  assert out.count("there"     ) == 10
  assert len(_lowerCAmelCase     ) == 20

 # check that we get items as fast as possible
 with Pool(2     ) as pool:
  UpperCamelCase	:					List[str]       	=   []
  for yield_time, content in iflatmap_unordered(
      _lowerCAmelCase							,   _aseconds_generator_of_aitems_with_timing							,   kwargs_iterable=[{"content": "a"}, {"content": "b"}]     ):
   assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
   out.append(_lowerCAmelCase     )
  assert out.count("a"     ) == 2
  assert out.count("b"     ) == 2
  assert len(_lowerCAmelCase     ) == 4

import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
   import jax
   import jax.numpy as jnp
   from flax.jax_utils import replicate
   from flax.training.common_utils import shard





@slow
@require_flax
class 		_lowercase						(					unittest.TestCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						str:

  """simple docstring"""





  super().tearDown()
  gc.collect()

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]			)   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_, UpperCamelCase_					:						str =					FlaxStableDiffusionPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2'   ,      revision='bf16'   ,      dtype=jnp.bfloataa   ,      )

  UpperCamelCase_					:						str =					'A painting of a squirrel eating a burger'
  UpperCamelCase_					:						Any =					jax.device_count()
  UpperCamelCase_					:						List[str] =					num_samples * [prompt]
  UpperCamelCase_					:						List[Any] =					sd_pipe.prepare_inputs(snake_case			)

  UpperCamelCase_					:						Dict =					replicate(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					shard(snake_case			)

  UpperCamelCase_					:						Dict =					jax.random.PRNGKey(0			)
  UpperCamelCase_					:						Tuple =					jax.random.split(snake_case   ,      jax.device_count()			)

  UpperCamelCase_					:						Optional[Any] =					sd_pipe(snake_case   ,      snake_case   ,      snake_case   ,      num_inference_steps=2_5   ,      jit=snake_case			)[0]
  assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

  UpperCamelCase_					:						int =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
  UpperCamelCase_					:						Tuple =					images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

  UpperCamelCase_					:						str =					jnp.asarray(jax.device_get(image_slice.flatten()			)			)
  UpperCamelCase_					:						List[str] =					jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]			)
  print(f"output_slice: {output_slice}"			)
  assert jnp.abs(output_slice - expected_slice			).max() < 1e-2


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						List[Any]:

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					'stabilityai/stable-diffusion-2'
  UpperCamelCase_, UpperCamelCase_					:						Tuple =					FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case   ,      subfolder='scheduler'			)
  UpperCamelCase_, UpperCamelCase_					:						Union[str, Any] =					FlaxStableDiffusionPipeline.from_pretrained(
      snake_case   ,      scheduler=snake_case   ,      revision='bf16'   ,      dtype=jnp.bfloataa   ,      )
  UpperCamelCase_					:						Optional[int] =					scheduler_params

  UpperCamelCase_					:						Optional[Any] =					'A painting of a squirrel eating a burger'
  UpperCamelCase_					:						Union[str, Any] =					jax.device_count()
  UpperCamelCase_					:						Union[str, Any] =					num_samples * [prompt]
  UpperCamelCase_					:						Tuple =					sd_pipe.prepare_inputs(snake_case			)

  UpperCamelCase_					:						List[Any] =					replicate(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					shard(snake_case			)

  UpperCamelCase_					:						Tuple =					jax.random.PRNGKey(0			)
  UpperCamelCase_					:						str =					jax.random.split(snake_case   ,      jax.device_count()			)

  UpperCamelCase_					:						str =					sd_pipe(snake_case   ,      snake_case   ,      snake_case   ,      num_inference_steps=2_5   ,      jit=snake_case			)[0]
  assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

  UpperCamelCase_					:						str =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
  UpperCamelCase_					:						int =					images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

  UpperCamelCase_					:						int =					jnp.asarray(jax.device_get(image_slice.flatten()			)			)
  UpperCamelCase_					:						Union[str, Any] =					jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]			)
  print(f"output_slice: {output_slice}"			)
  assert jnp.abs(output_slice - expected_slice			).max() < 1e-2



