






import copy
import unittest

from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
     MODEL_FOR_QUESTION_ANSWERING_MAPPING,
     MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
     MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
     LayoutLMvaConfig,
     LayoutLMvaForQuestionAnswering,
     LayoutLMvaForSequenceClassification,
     LayoutLMvaForTokenClassification,
     LayoutLMvaModel,
 )
 from transformers.models.layoutlmva.modeling_layoutlmva import LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import LayoutLMvaImageProcessor





class 			__a  :







 def __init__(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   List[Any]=2							,   SCREAMING_SNAKE_CASE					:   Optional[int]=3							,   SCREAMING_SNAKE_CASE					:   Optional[int]=4							,   SCREAMING_SNAKE_CASE					:   Any=2							,   SCREAMING_SNAKE_CASE					:   str=7							,   SCREAMING_SNAKE_CASE					:   Optional[Any]=True							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=True							,   SCREAMING_SNAKE_CASE					:   Optional[int]=True							,   SCREAMING_SNAKE_CASE					:   Optional[Any]=True							,   SCREAMING_SNAKE_CASE					:   Optional[int]=99							,   SCREAMING_SNAKE_CASE					:   List[str]=36							,   SCREAMING_SNAKE_CASE					:   int=3							,   SCREAMING_SNAKE_CASE					:   Optional[Any]=4							,   SCREAMING_SNAKE_CASE					:   Dict=37							,   SCREAMING_SNAKE_CASE					:   Optional[int]="gelu"							,   SCREAMING_SNAKE_CASE					:   Tuple=0.1							,   SCREAMING_SNAKE_CASE					:   Optional[Any]=0.1							,   SCREAMING_SNAKE_CASE					:   Any=5_12							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=16							,   SCREAMING_SNAKE_CASE					:   List[str]=2							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=0.0_2							,   SCREAMING_SNAKE_CASE					:   Any=6							,   SCREAMING_SNAKE_CASE					:   List[str]=6							,   SCREAMING_SNAKE_CASE					:   int=3							,   SCREAMING_SNAKE_CASE					:   Tuple=4							,   SCREAMING_SNAKE_CASE					:   Tuple=None							,   SCREAMING_SNAKE_CASE					:   Optional[int]=10_00							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  parent
      UpperCamelCase__			:   Union[str, Any]       				=  batch_size
      UpperCamelCase__			:   List[Any]       				=  num_channels
      UpperCamelCase__			:   List[Any]       				=  image_size
      UpperCamelCase__			:   Any       				=  patch_size
      UpperCamelCase__			:   int       				=  text_seq_length
      UpperCamelCase__			:   List[Any]       				=  is_training
      UpperCamelCase__			:   str       				=  use_input_mask
      UpperCamelCase__			:   Dict       				=  use_token_type_ids
      UpperCamelCase__			:   Dict       				=  use_labels
      UpperCamelCase__			:   Dict       				=  vocab_size
      UpperCamelCase__			:   Dict       				=  hidden_size
      UpperCamelCase__			:   Dict       				=  num_hidden_layers
      UpperCamelCase__			:   List[str]       				=  num_attention_heads
      UpperCamelCase__			:   str       				=  intermediate_size
      UpperCamelCase__			:   Optional[int]       				=  hidden_act
      UpperCamelCase__			:   Optional[Any]       				=  hidden_dropout_prob
      UpperCamelCase__			:   Union[str, Any]       				=  attention_probs_dropout_prob
      UpperCamelCase__			:   List[Any]       				=  max_position_embeddings
      UpperCamelCase__			:   int       				=  type_vocab_size
      UpperCamelCase__			:   List[Any]       				=  type_sequence_label_size
      UpperCamelCase__			:   Union[str, Any]       				=  initializer_range
      UpperCamelCase__			:   List[Any]       				=  coordinate_size
      UpperCamelCase__			:   Any       				=  shape_size
      UpperCamelCase__			:   Tuple       				=  num_labels
      UpperCamelCase__			:   int       				=  num_choices
      UpperCamelCase__			:   Any       				=  scope
      UpperCamelCase__			:   Optional[Any]       				=  range_bbox

      # LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
      UpperCamelCase__			:   int       				=  text_seq_length
      UpperCamelCase__			:   Union[str, Any]       				=  (image_size // patch_size) ** 2 + 1
      UpperCamelCase__			:   Tuple       				=  self.text_seq_length + self.image_seq_length







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  ids_tensor([self.batch_size, self.text_seq_length]							,   self.vocab_size						)

      UpperCamelCase__			:   str       				=  ids_tensor([self.batch_size, self.text_seq_length, 4]							,   self.range_bbox						)
      # Ensure that bbox is legal
      for i in range(bbox.shape[0]						):
           for j in range(bbox.shape[1]						):
                if bbox[i, j, 3] < bbox[i, j, 1]:
                     UpperCamelCase__			:   Optional[Any]       				=  bbox[i, j, 3]
                     UpperCamelCase__			:   Dict       				=  bbox[i, j, 1]
                     UpperCamelCase__			:   str       				=  t
                if bbox[i, j, 2] < bbox[i, j, 0]:
                     UpperCamelCase__			:   Dict       				=  bbox[i, j, 2]
                     UpperCamelCase__			:   Any       				=  bbox[i, j, 0]
                     UpperCamelCase__			:   Optional[int]       				=  t

      UpperCamelCase__			:   int       				=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

      UpperCamelCase__			:   str       				=  None
      if self.use_input_mask:
           UpperCamelCase__			:   str       				=  random_attention_mask([self.batch_size, self.text_seq_length]						)

      UpperCamelCase__			:   str       				=  None
      if self.use_token_type_ids:
           UpperCamelCase__			:   List[str]       				=  ids_tensor([self.batch_size, self.text_seq_length]							,   self.type_vocab_size						)

      UpperCamelCase__			:   List[str]       				=  None
      UpperCamelCase__			:   Union[str, Any]       				=  None
      if self.use_labels:
           UpperCamelCase__			:   Tuple       				=  ids_tensor([self.batch_size]							,   self.type_sequence_label_size						)
           UpperCamelCase__			:   List[Any]       				=  ids_tensor([self.batch_size, self.text_seq_length]							,   self.num_labels						)

      UpperCamelCase__			:   Dict       				=  LayoutLMvaConfig(
          vocab_size=self.vocab_size							,   hidden_size=self.hidden_size							,   num_hidden_layers=self.num_hidden_layers							,   num_attention_heads=self.num_attention_heads							,   intermediate_size=self.intermediate_size							,   hidden_act=self.hidden_act							,   hidden_dropout_prob=self.hidden_dropout_prob							,   attention_probs_dropout_prob=self.attention_probs_dropout_prob							,   max_position_embeddings=self.max_position_embeddings							,   type_vocab_size=self.type_vocab_size							,   initializer_range=self.initializer_range							,   coordinate_size=self.coordinate_size							,   shape_size=self.shape_size							,   input_size=self.image_size							,   patch_size=self.patch_size							,   )

      return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  LayoutLMvaModel(config=_lowerCamelCase						)
      model.to(_lowerCamelCase						)
      model.eval()

      # text + image
      UpperCamelCase__			:   str       				=  model(_lowerCamelCase							,   pixel_values=_lowerCamelCase						)
      UpperCamelCase__			:   int       				=  model(
          _lowerCamelCase							,   bbox=_lowerCamelCase							,   pixel_values=_lowerCamelCase							,   attention_mask=_lowerCamelCase							,   token_type_ids=_lowerCamelCase						)
      UpperCamelCase__			:   List[str]       				=  model(_lowerCamelCase							,   bbox=_lowerCamelCase							,   pixel_values=_lowerCamelCase							,   token_type_ids=_lowerCamelCase						)
      UpperCamelCase__			:   Tuple       				=  model(_lowerCamelCase							,   bbox=_lowerCamelCase							,   pixel_values=_lowerCamelCase						)

      self.parent.assertEqual(result.last_hidden_state.shape							,   (self.batch_size, self.seq_length, self.hidden_size)						)

      # text only
      UpperCamelCase__			:   Union[str, Any]       				=  model(_lowerCamelCase						)
      self.parent.assertEqual(
          result.last_hidden_state.shape							,   (self.batch_size, self.text_seq_length, self.hidden_size)						)

      # image only
      UpperCamelCase__			:   Any       				=  model(pixel_values=_lowerCamelCase						)
      self.parent.assertEqual(
          result.last_hidden_state.shape							,   (self.batch_size, self.image_seq_length, self.hidden_size)						)







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  self.num_labels
      UpperCamelCase__			:   Dict       				=  LayoutLMvaForSequenceClassification(_lowerCamelCase						)
      model.to(_lowerCamelCase						)
      model.eval()
      UpperCamelCase__			:   Union[str, Any]       				=  model(
          _lowerCamelCase							,   bbox=_lowerCamelCase							,   pixel_values=_lowerCamelCase							,   attention_mask=_lowerCamelCase							,   token_type_ids=_lowerCamelCase							,   labels=_lowerCamelCase							,   )
      self.parent.assertEqual(result.logits.shape							,   (self.batch_size, self.num_labels)						)







 def 					__lowercase					(		self					:   str							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  self.num_labels
      UpperCamelCase__			:   Optional[Any]       				=  LayoutLMvaForTokenClassification(config=_lowerCamelCase						)
      model.to(_lowerCamelCase						)
      model.eval()
      UpperCamelCase__			:   Any       				=  model(
          _lowerCamelCase							,   bbox=_lowerCamelCase							,   pixel_values=_lowerCamelCase							,   attention_mask=_lowerCamelCase							,   token_type_ids=_lowerCamelCase							,   labels=_lowerCamelCase							,   )
      self.parent.assertEqual(result.logits.shape							,   (self.batch_size, self.text_seq_length, self.num_labels)						)







 def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  LayoutLMvaForQuestionAnswering(config=_lowerCamelCase						)
      model.to(_lowerCamelCase						)
      model.eval()
      UpperCamelCase__			:   Dict       				=  model(
          _lowerCamelCase							,   bbox=_lowerCamelCase							,   pixel_values=_lowerCamelCase							,   attention_mask=_lowerCamelCase							,   token_type_ids=_lowerCamelCase							,   start_positions=_lowerCamelCase							,   end_positions=_lowerCamelCase							,   )
      self.parent.assertEqual(result.start_logits.shape							,   (self.batch_size, self.seq_length)						)
      self.parent.assertEqual(result.end_logits.shape							,   (self.batch_size, self.seq_length)						)







 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  self.prepare_config_and_inputs()
      (
          (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,     (
          UpperCamelCase__     
      )				,          
      )			:   Optional[Any]       				=  config_and_inputs
      UpperCamelCase__			:   Tuple       				=  {
          "input_ids": input_ids,
          "bbox": bbox,
          "pixel_values": pixel_values,
          "token_type_ids": token_type_ids,
          "attention_mask": input_mask,
      }
      return config, inputs_dict







@require_torch
class 			__a  (    a_       ,   a_       ,   unittest.TestCase       ):
 _lowerCAmelCase       :       List[Any]        = False
 _lowerCAmelCase       :       Any        = False
 _lowerCAmelCase       :       Optional[int]        = False

 _lowerCAmelCase       :       Union[str, Any]        = (
     (
         LayoutLMvaModel,
         LayoutLMvaForSequenceClassification,
         LayoutLMvaForTokenClassification,
         LayoutLMvaForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 _lowerCAmelCase       :       Dict        = (
     {"document-question-answering": LayoutLMvaForQuestionAnswering, "feature-extraction": LayoutLMvaModel}
     if is_torch_available()
     else {}
 )







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]						):


      '''simple docstring'''






      return True







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  LayoutLMvaModelTester(self						)
      UpperCamelCase__			:   Tuple       				=  ConfigTester(self							,   config_class=_lowerCamelCase							,   hidden_size=37						)







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=False						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  copy.deepcopy(_lowerCamelCase						)
      if model_class in get_values(_lowerCamelCase						):
           UpperCamelCase__			:   Union[str, Any]       				=  {
               k: v.unsqueeze(1						).expand(-1							,   self.model_tester.num_choices							,   -1						).contiguous()
               if isinstance(_lowerCamelCase							,   torch.Tensor						) and v.ndim > 1
               else v
               for k, v in inputs_dict.items()
           }
      if return_labels:
           if model_class in get_values(_lowerCamelCase						):
                UpperCamelCase__			:   Optional[int]       				=  torch.ones(self.model_tester.batch_size							,   dtype=torch.long							,   device=_lowerCamelCase						)
           elif model_class in get_values(_lowerCamelCase						):
                UpperCamelCase__			:   Optional[Any]       				=  torch.zeros(
                    self.model_tester.batch_size							,   dtype=torch.long							,   device=_lowerCamelCase						)
                UpperCamelCase__			:   Any       				=  torch.zeros(
                    self.model_tester.batch_size							,   dtype=torch.long							,   device=_lowerCamelCase						)
           elif model_class in [
               *get_values(_lowerCamelCase						),
           ]:
                UpperCamelCase__			:   Union[str, Any]       				=  torch.zeros(
                    self.model_tester.batch_size							,   dtype=torch.long							,   device=_lowerCamelCase						)
           elif model_class in [
               *get_values(_lowerCamelCase						),
           ]:
                UpperCamelCase__			:   List[str]       				=  torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.text_seq_length)							,   dtype=torch.long							,   device=_lowerCamelCase							,   )

      return inputs_dict







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      self.config_tester.run_common_tests()







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_lowerCamelCase						)







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
           UpperCamelCase__			:   Tuple       				=  type
           self.model_tester.create_and_check_model(*_lowerCamelCase						)







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Any       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase						)







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase						)







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase						)







 @slow
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      for model_name in LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase__			:   str       				=  LayoutLMvaModel.from_pretrained(_lowerCamelCase						)
           self.assertIsNotNone(_lowerCamelCase						)







def 		SCREAMING_SNAKE_CASE							(					) ->	str:
     UpperCamelCase__			:   Union[str, Any]       				=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
     return image





@require_torch
class 			__a  (    unittest.TestCase       ):







 @cached_property
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      return LayoutLMvaImageProcessor(apply_ocr=_lowerCamelCase						) if is_vision_available() else None







 @slow
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  LayoutLMvaModel.from_pretrained("microsoft/layoutlmv3-base"						).to(_lowerCamelCase						)

      UpperCamelCase__			:   List[str]       				=  self.default_image_processor
      UpperCamelCase__			:   List[Any]       				=  prepare_img()
      UpperCamelCase__			:   int       				=  image_processor(images=_lowerCamelCase							,   return_tensors="pt"						).pixel_values.to(_lowerCamelCase						)

      UpperCamelCase__			:   Dict       				=  torch.tensor([[1, 2]]						)
      UpperCamelCase__			:   Optional[Any]       				=  torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]						).unsqueeze(0						)

      # forward pass
      UpperCamelCase__			:   Tuple       				=  model(
          input_ids=input_ids.to(_lowerCamelCase						)							,   bbox=bbox.to(_lowerCamelCase						)							,   pixel_values=pixel_values.to(_lowerCamelCase						)							,   )

      # verify the logits
      UpperCamelCase__			:   str       				=  torch.Size((1, 1_99, 7_68)						)
      self.assertEqual(outputs.last_hidden_state.shape							,   _lowerCamelCase						)

      UpperCamelCase__			:   int       				=  torch.tensor(
          [[-0.0_5_2_9, 0.3_6_1_8, 0.1_6_3_2], [-0.1_5_8_7, -0.1_6_6_7, -0.0_4_0_0], [-0.1_5_5_7, -0.1_6_7_1, -0.0_5_0_5]]						).to(_lowerCamelCase						)

      self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]							,   _lowerCamelCase							,   atol=1e-4						)						)







"""simple docstring"""



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        a     (   unittest.TestCase	):







    def 			UpperCamelCase_      (			self       ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()







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

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       sd_pipe(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		num_inference_steps=2_5							,		jit=_lowerCamelCase       )[0]
      assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

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

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       'stabilityai/stable-diffusion-2'
      lowercase    ,     lowercase							    =       FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCamelCase							,		subfolder='scheduler'       )
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          _lowerCamelCase							,		scheduler=_lowerCamelCase							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)
      lowercase							    =       scheduler_params

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       sd_pipe(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		num_inference_steps=2_5							,		jit=_lowerCamelCase       )[0]
      assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

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

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2


