"""simple docstring"""




import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

			if is_vision_available():
						from transformers import MaskaFormerImageProcessor

if is_vision_available():
			from PIL import Image

class    a_   :




					def __init__(						self    :  int	,    snake_case__    :  List[str]	,    snake_case__    :  Optional[Any]=2	,    snake_case__    :  str=True	,    snake_case__    :  List[str]=False	,    snake_case__    :  Optional[int]=10	,    snake_case__    :  List[str]=3	,    snake_case__    :  Tuple=32 * 8	,    snake_case__    :  Dict=32 * 8	,    snake_case__    :  Union[str, Any]=4	,    snake_case__    :  Union[str, Any]=64	,    ):
							lowerCAmelCase__   = parent
							lowerCAmelCase__   = batch_size
							lowerCAmelCase__   = is_training
							lowerCAmelCase__   = use_auxiliary_loss
							lowerCAmelCase__   = num_queries
							lowerCAmelCase__   = num_channels
							lowerCAmelCase__   = min_size
							lowerCAmelCase__   = max_size
							lowerCAmelCase__   = num_labels
							lowerCAmelCase__   = hidden_dim
							lowerCAmelCase__   = hidden_dim




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]  ):
							lowerCAmelCase__   = floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]  ).to(
							    snake_case__  )

							lowerCAmelCase__   = torch.ones([self.batch_size, self.min_size, self.max_size]	,    device=snake_case__  )

							lowerCAmelCase__   = (
							    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]	,    device=snake_case__  ) > 0.5
							).float()
							lowerCAmelCase__   = (torch.rand((self.batch_size, self.num_labels)	,    device=snake_case__  ) > 0.5).long()

							lowerCAmelCase__   = self.get_config()
							return config, pixel_values, pixel_mask, mask_labels, class_labels




					def     _SCREAMING_SNAKE_CASE    (						self    :  Union[str, Any]  ):
							lowerCAmelCase__   = MaskaFormerConfig(
							    hidden_size=self.hidden_dim	,    )
							lowerCAmelCase__   = self.num_queries
							lowerCAmelCase__   = self.num_labels

							lowerCAmelCase__   = [1, 1, 1, 1]
							lowerCAmelCase__   = self.num_channels

							lowerCAmelCase__   = 64
							lowerCAmelCase__   = 128
							lowerCAmelCase__   = self.hidden_dim
							lowerCAmelCase__   = self.hidden_dim
							lowerCAmelCase__   = self.hidden_dim
							return config




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[Any]  ):
							lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__   = self.prepare_config_and_inputs()
							lowerCAmelCase__   = {"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
							return config, inputs_dict




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[Any]	,    snake_case__    :  Union[str, Any]	,    snake_case__    :  Optional[Any]  ):
							lowerCAmelCase__   = output.encoder_hidden_states
							lowerCAmelCase__   = output.pixel_decoder_hidden_states
							lowerCAmelCase__   = output.transformer_decoder_hidden_states

							self.parent.assertTrue(len(snake_case__  )	,    len(config.backbone_config.depths  )  )
							self.parent.assertTrue(len(snake_case__  )	,    len(config.backbone_config.depths  )  )
							self.parent.assertTrue(len(snake_case__  )	,    config.decoder_layers  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Any	,    snake_case__    :  List[str]	,    snake_case__    :  List[str]	,    snake_case__    :  Optional[int]	,    snake_case__    :  Union[str, Any]=False  ):
							with torch.no_grad():
									lowerCAmelCase__   = MaskaFormerModel(config=snake_case__  )
									model.to(snake_case__  )
									model.eval()

									lowerCAmelCase__   = model(pixel_values=snake_case__	,    pixel_mask=snake_case__  )
									lowerCAmelCase__   = model(snake_case__	,    output_hidden_states=snake_case__  )

							self.parent.assertEqual(
							    output.transformer_decoder_last_hidden_state.shape	,    (self.batch_size, self.num_queries, self.hidden_dim)	,    )
							# let's ensure the other two hidden state exists
							self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None  )
							self.parent.assertTrue(output.encoder_last_hidden_state is not None  )

							if output_hidden_states:
									self.check_output_hidden_state(snake_case__	,    snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  str	,    snake_case__    :  List[str]	,    snake_case__    :  Optional[int]	,    snake_case__    :  Tuple	,    snake_case__    :  List[str]	,    snake_case__    :  int  ):
							lowerCAmelCase__   = MaskaFormerForUniversalSegmentation(config=snake_case__  )
							model.to(snake_case__  )
							model.eval()

							def comm_check_on_output(snake_case__    :  Dict  ):
									# let's still check that all the required stuff is there
									self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None  )
									self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None  )
									self.parent.assertTrue(result.encoder_last_hidden_state is not None  )
									# okay, now we need to check the logits shape
									# due to the encoder compression, masks have a //4 spatial size
									self.parent.assertEqual(
									    result.masks_queries_logits.shape	,    (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)	,    )
									# + 1 for null class
									self.parent.assertEqual(
									    result.class_queries_logits.shape	,    (self.batch_size, self.num_queries, self.num_labels + 1)  )

							with torch.no_grad():
									lowerCAmelCase__   = model(pixel_values=snake_case__	,    pixel_mask=snake_case__  )
									lowerCAmelCase__   = model(snake_case__  )

									comm_check_on_output(snake_case__  )

									lowerCAmelCase__   = model(
									    pixel_values=snake_case__	,    pixel_mask=snake_case__	,    mask_labels=snake_case__	,    class_labels=snake_case__  )

							comm_check_on_output(snake_case__  )

							self.parent.assertTrue(result.loss is not None  )
							self.parent.assertEqual(result.loss.shape	,    torch.Size([1]  )  )

@require_torch
class    a_   (  __UpperCamelCase			,		__UpperCamelCase			,		unittest.TestCase      ):
					UpperCamelCase_			:						Tuple      =     (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
					UpperCamelCase_			:						Union[str, Any]      =     {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

					UpperCamelCase_			:						int      =     False
					UpperCamelCase_			:						Tuple      =     False
					UpperCamelCase_			:						Optional[int]      =     False
					UpperCamelCase_			:						str      =     False




					def     _SCREAMING_SNAKE_CASE    (						self    :  str  ):
							lowerCAmelCase__   = MaskaFormerModelTester(self  )
							lowerCAmelCase__   = ConfigTester(self	,    config_class=snake_case__	,    has_text_modality=snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  int  ):
							self.config_tester.run_common_tests()




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[int]  ):
							lowerCAmelCase__ ,       lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.create_and_check_maskaformer_model(snake_case__	,    **snake_case__	,    output_hidden_states=snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]  ):
							lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*snake_case__  )




					@unittest.skip(reason="""Mask2Former does not use inputs_embeds"""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  str  ):
							pass




					@unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							pass




					@unittest.skip(reason="""Mask2Former is not a generative model"""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							pass




					@unittest.skip(reason="""Mask2Former does not use token embeddings"""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							pass




					@require_torch_multi_gpu
					@unittest.skip(
					    reason="""Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  int  ):
							pass




					@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  int  ):
							pass




					def     _SCREAMING_SNAKE_CASE    (						self    :  Union[str, Any]  ):
							lowerCAmelCase__ ,       lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
									lowerCAmelCase__   = model_class(snake_case__  )
									lowerCAmelCase__   = inspect.signature(model.forward  )
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									lowerCAmelCase__   = [*signature.parameters.keys()]

									lowerCAmelCase__   = ["""pixel_values"""]
									self.assertListEqual(arg_names[:1]	,    snake_case__  )




					@slow
					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
									lowerCAmelCase__   = MaskaFormerModel.from_pretrained(snake_case__  )
									self.assertIsNotNone(snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]  ):
							lowerCAmelCase__   = (self.model_tester.min_size,) * 2
							lowerCAmelCase__   = {
							    """pixel_values""": torch.randn((2, 3, *size)	,    device=snake_case__  ),
							    """mask_labels""": torch.randn((2, 10, *size)	,    device=snake_case__  ),
							    """class_labels""": torch.zeros(2	,    10	,    device=snake_case__  ).long(),
							}
							lowerCAmelCase__   = self.model_tester.get_config()

							lowerCAmelCase__   = MaskaFormerForUniversalSegmentation(snake_case__  ).to(snake_case__  )
							lowerCAmelCase__   = model(**snake_case__  )
							self.assertTrue(outputs.loss is not None  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  int  ):
							lowerCAmelCase__ ,       lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.create_and_check_maskaformer_model(snake_case__	,    **snake_case__	,    output_hidden_states=snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[Any]  ):
							lowerCAmelCase__ ,       lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
									lowerCAmelCase__   = model_class(snake_case__  ).to(snake_case__  )
									lowerCAmelCase__   = model(**snake_case__	,    output_attentions=snake_case__  )
									self.assertTrue(outputs.attentions is not None  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							if not self.model_tester.is_training:
									return

							lowerCAmelCase__   = self.all_model_classes[1]
							lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs()

							lowerCAmelCase__   = model_class(snake_case__  )
							model.to(snake_case__  )
							model.train()

							lowerCAmelCase__   = model(snake_case__	,    mask_labels=snake_case__	,    class_labels=snake_case__  ).loss
							loss.backward()




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[int]  ):
							lowerCAmelCase__   = self.all_model_classes[1]
							lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs()
							lowerCAmelCase__   = True
							lowerCAmelCase__   = True

							lowerCAmelCase__   = model_class(snake_case__  ).to(snake_case__  )
							model.train()

							lowerCAmelCase__   = model(snake_case__	,    mask_labels=snake_case__	,    class_labels=snake_case__  )

							lowerCAmelCase__   = outputs.encoder_hidden_states[0]
							encoder_hidden_states.retain_grad()

							lowerCAmelCase__   = outputs.pixel_decoder_hidden_states[0]
							pixel_decoder_hidden_states.retain_grad()

							lowerCAmelCase__   = outputs.transformer_decoder_hidden_states[0]
							transformer_decoder_hidden_states.retain_grad()

							lowerCAmelCase__   = outputs.attentions[0]
							attentions.retain_grad()

							outputs.loss.backward(retain_graph=snake_case__  )

							self.assertIsNotNone(encoder_hidden_states.grad  )
							self.assertIsNotNone(pixel_decoder_hidden_states.grad  )
							self.assertIsNotNone(transformer_decoder_hidden_states.grad  )
							self.assertIsNotNone(attentions.grad  )
__lowerCAmelCase     :				int    				=    1e-4







def 				_UpperCAmelCase		(							):





		"""simple docstring"""




		lowerCAmelCase__   = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""  )
		return image






@require_vision
@slow
class    a_   (  unittest.TestCase      ):




					@cached_property
					def     _SCREAMING_SNAKE_CASE    (						self    :  Tuple  ):
							return "facebook/mask2former-swin-small-coco-instance"




					@cached_property
					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]  ):
							return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints  ) if is_vision_available() else None




					def     _SCREAMING_SNAKE_CASE    (						self    :  Union[str, Any]  ):
							lowerCAmelCase__   = MaskaFormerModel.from_pretrained(self.model_checkpoints  ).to(snake_case__  )
							lowerCAmelCase__   = self.default_image_processor
							lowerCAmelCase__   = prepare_img()
							lowerCAmelCase__   = image_processor(snake_case__	,    return_tensors="""pt"""  ).to(snake_case__  )
							lowerCAmelCase__   = inputs["""pixel_values"""].shape
							# check size is divisible by 32
							self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0  )
							# check size
							self.assertEqual(snake_case__	,    (1, 3, 384, 384)  )

							with torch.no_grad():
									lowerCAmelCase__   = model(**snake_case__  )

							lowerCAmelCase__   = torch.tensor(
							    [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]  ).to(snake_case__  )
							self.assertTrue(
							    torch.allclose(
							        outputs.encoder_last_hidden_state[0, 0, :3, :3]	,    snake_case__	,    atol=snake_case__  )  )

							lowerCAmelCase__   = torch.tensor(
							    [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]  ).to(snake_case__  )
							self.assertTrue(
							    torch.allclose(
							        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]	,    snake_case__	,    atol=snake_case__  )  )

							lowerCAmelCase__   = torch.tensor(
							    [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]  ).to(snake_case__  )
							self.assertTrue(
							    torch.allclose(
							        outputs.transformer_decoder_last_hidden_state[0, :3, :3]	,    snake_case__	,    atol=snake_case__  )  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  int  ):
							lowerCAmelCase__   = MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints  ).to(snake_case__  ).eval()
							lowerCAmelCase__   = self.default_image_processor
							lowerCAmelCase__   = prepare_img()
							lowerCAmelCase__   = image_processor(snake_case__	,    return_tensors="""pt"""  ).to(snake_case__  )
							lowerCAmelCase__   = inputs["""pixel_values"""].shape
							# check size is divisible by 32
							self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0  )
							# check size
							self.assertEqual(snake_case__	,    (1, 3, 384, 384)  )

							with torch.no_grad():
									lowerCAmelCase__   = model(**snake_case__  )
							# masks_queries_logits
							lowerCAmelCase__   = outputs.masks_queries_logits
							self.assertEqual(
							    masks_queries_logits.shape	,    (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)  )
							lowerCAmelCase__   = [
							    [-8.7839, -9.0056, -8.8121],
							    [-7.4104, -7.0313, -6.5401],
							    [-6.6105, -6.3427, -6.4675],
							]
							lowerCAmelCase__   = torch.tensor(snake_case__  ).to(snake_case__  )
							self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]	,    snake_case__	,    atol=snake_case__  )  )
							# class_queries_logits
							lowerCAmelCase__   = outputs.class_queries_logits
							self.assertEqual(class_queries_logits.shape	,    (1, model.config.num_queries, model.config.num_labels + 1)  )
							lowerCAmelCase__   = torch.tensor(
							    [
							        [1.8324, -8.0835, -4.1922],
							        [0.8450, -9.0050, -3.6053],
							        [0.3045, -7.7293, -3.0275],
							    ]  ).to(snake_case__  )
							self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]	,    snake_case__	,    atol=snake_case__  )  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[Any]  ):
							lowerCAmelCase__   = MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints  ).to(snake_case__  ).eval()
							lowerCAmelCase__   = self.default_image_processor

							lowerCAmelCase__   = image_processor(
							    [np.zeros((3, 800, 1333)  ), np.zeros((3, 800, 1333)  )]	,    segmentation_maps=[np.zeros((384, 384)  ).astype(np.floataa  ), np.zeros((384, 384)  ).astype(np.floataa  )]	,    return_tensors="""pt"""	,    )

							lowerCAmelCase__   = inputs["""pixel_values"""].to(snake_case__  )
							lowerCAmelCase__   = [el.to(snake_case__  ) for el in inputs["""mask_labels"""]]
							lowerCAmelCase__   = [el.to(snake_case__  ) for el in inputs["""class_labels"""]]

							with torch.no_grad():
									lowerCAmelCase__   = model(**snake_case__  )

							self.assertTrue(outputs.loss is not None  )





"""simple docstring"""




from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule







def 				_UpperCAmelCase		(							lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__=1e-12  ):





		"""simple docstring"""




		lowerCAmelCase__   = jnp.divide(emb_a.T      ,	jnp.clip(jnp.linalg.norm(lowerCamelCase__      ,	axis=1  )      ,	a_min=lowerCamelCase__  )  ).T
		lowerCAmelCase__   = jnp.divide(emb_a.T      ,	jnp.clip(jnp.linalg.norm(lowerCamelCase__      ,	axis=1  )      ,	a_min=lowerCamelCase__  )  ).T
		return jnp.matmul(lowerCamelCase__      ,	norm_emb_a.T  )

class    a_   (  nn.Module      ):
					UpperCamelCase_			:						CLIPConfig
					UpperCamelCase_			:						jnp.dtype      =     jnp.floataa




					def     _SCREAMING_SNAKE_CASE    (						self    :  Tuple  ):
							lowerCAmelCase__   = FlaxCLIPVisionModule(self.config.vision_config  )
							lowerCAmelCase__   = nn.Dense(self.config.projection_dim	,    use_bias=snake_case__	,    dtype=self.dtype  )

							lowerCAmelCase__   = self.param("""concept_embeds"""	,    jax.nn.initializers.ones	,    (17, self.config.projection_dim)  )
							lowerCAmelCase__   = self.param(
							    """special_care_embeds"""	,    jax.nn.initializers.ones	,    (3, self.config.projection_dim)  )

							lowerCAmelCase__   = self.param("""concept_embeds_weights"""	,    jax.nn.initializers.ones	,    (17,)  )
							lowerCAmelCase__   = self.param("""special_care_embeds_weights"""	,    jax.nn.initializers.ones	,    (3,)  )




					def __call__(						self    :  Optional[int]	,    snake_case__    :  int  ):
							lowerCAmelCase__   = self.vision_model(snake_case__  )[1]
							lowerCAmelCase__   = self.visual_projection(snake_case__  )

							lowerCAmelCase__   = jax_cosine_distance(snake_case__	,    self.special_care_embeds  )
							lowerCAmelCase__   = jax_cosine_distance(snake_case__	,    self.concept_embeds  )

							# increase this value to create a stronger `nfsw` filter
							# at the cost of increasing the possibility of filtering benign image inputs
							lowerCAmelCase__   = 0.0

							lowerCAmelCase__   = special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
							lowerCAmelCase__   = jnp.round(snake_case__	,    3  )
							lowerCAmelCase__   = jnp.any(special_scores > 0	,    axis=1	,    keepdims=snake_case__  )
							# Use a lower threshold if an image has any special care concept
							lowerCAmelCase__   = is_special_care * 0.01

							lowerCAmelCase__   = cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
							lowerCAmelCase__   = jnp.round(snake_case__	,    3  )
							lowerCAmelCase__   = jnp.any(concept_scores > 0	,    axis=1  )

							return has_nsfw_concepts

class    a_   (  __UpperCamelCase      ):
					UpperCamelCase_			:						Any      =     CLIPConfig
					UpperCamelCase_			:						str      =     "clip_input"
					UpperCamelCase_			:						str      =     FlaxStableDiffusionSafetyCheckerModule




					def __init__(						self    :  Optional[int]	,    snake_case__    :  CLIPConfig	,    snake_case__    :  Optional[Tuple] = None	,    snake_case__    :  int = 0	,    snake_case__    :  jnp.dtype = jnp.floataa	,    snake_case__    :  bool = True	,    **snake_case__    :  str	,    ):
							if input_shape is None:
									lowerCAmelCase__   = (1, 224, 224, 3)
							lowerCAmelCase__   = self.module_class(config=snake_case__	,    dtype=snake_case__	,    **snake_case__  )
							super().__init__(snake_case__	,    snake_case__	,    input_shape=snake_case__	,    seed=snake_case__	,    dtype=snake_case__	,    _do_init=_do_init  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Tuple	,    snake_case__    :  jax.random.KeyArray	,    snake_case__    :  Tuple	,    snake_case__    :  FrozenDict = None  ):
							# init input tensor
							lowerCAmelCase__   = jax.random.normal(snake_case__	,    snake_case__  )

							lowerCAmelCase__ ,       lowerCAmelCase__   = jax.random.split(snake_case__  )
							lowerCAmelCase__   = {"""params""": params_rng, """dropout""": dropout_rng}

							lowerCAmelCase__   = self.module.init(snake_case__	,    snake_case__  )["""params"""]

							return random_params




					def __call__(						self    :  Union[str, Any]	,    snake_case__    :  Union[str, Any]	,    snake_case__    :  dict = None	,    ):
							lowerCAmelCase__   = jnp.transpose(snake_case__	,    (0, 2, 3, 1)  )

							return self.module.apply(
							    {"""params""": params or self.params}	,    jnp.array(snake_case__	,    dtype=jnp.floataa  )	,    rngs={}	,    )




