






import collections
import inspect
import unittest

from transformers import SwinvaConfig
from transformers.testing_utils import require_torch, require_vision, 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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import SwinvaForImageClassification, SwinvaForMaskedImageModeling, SwinvaModel
 from transformers.models.swinva.modeling_swinva import SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor


class 					_snake_case						:






 def __init__(		self					,UpperCamelCase					,UpperCamelCase=13					,UpperCamelCase=32					,UpperCamelCase=2					,UpperCamelCase=3					,UpperCamelCase=16					,UpperCamelCase=[1, 2, 1]					,UpperCamelCase=[2, 2, 4]					,UpperCamelCase=2					,UpperCamelCase=2.0					,UpperCamelCase=True					,UpperCamelCase=0.0					,UpperCamelCase=0.0					,UpperCamelCase=0.1					,UpperCamelCase="gelu"					,UpperCamelCase=False					,UpperCamelCase=True					,UpperCamelCase=0.02					,UpperCamelCase=1E-5					,UpperCamelCase=True					,UpperCamelCase=None					,UpperCamelCase=True					,UpperCamelCase=10					,UpperCamelCase=8					,)						->    int:
     snake_case__      :Optional[Any]         =   parent
     snake_case__      :int         =   batch_size
     snake_case__      :int         =   image_size
     snake_case__      :Union[str, Any]         =   patch_size
     snake_case__      :Tuple         =   num_channels
     snake_case__      :Any         =   embed_dim
     snake_case__      :int         =   depths
     snake_case__      :List[str]         =   num_heads
     snake_case__      :Union[str, Any]         =   window_size
     snake_case__      :Optional[Any]         =   mlp_ratio
     snake_case__      :List[Any]         =   qkv_bias
     snake_case__      :Union[str, Any]         =   hidden_dropout_prob
     snake_case__      :Dict         =   attention_probs_dropout_prob
     snake_case__      :List[str]         =   drop_path_rate
     snake_case__      :List[Any]         =   hidden_act
     snake_case__      :Optional[Any]         =   use_absolute_embeddings
     snake_case__      :Any         =   patch_norm
     snake_case__      :Union[str, Any]         =   layer_norm_eps
     snake_case__      :List[Any]         =   initializer_range
     snake_case__      :Any         =   is_training
     snake_case__      :List[Any]         =   scope
     snake_case__      :Optional[Any]         =   use_labels
     snake_case__      :Optional[Any]         =   type_sequence_label_size
     snake_case__      :Optional[Any]         =   encoder_stride






 def    lowerCAmelCase_							(		self       )						->    List[str]:
     snake_case__      :Optional[int]         =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

     snake_case__      :int         =   None
     if self.use_labels:
         snake_case__      :Dict         =   ids_tensor([self.batch_size]					,self.type_sequence_label_size       )

     snake_case__      :Tuple         =   self.get_config()

     return config, pixel_values, labels






 def    lowerCAmelCase_							(		self       )						->    int:
     return SwinvaConfig(
         image_size=self.image_size					,patch_size=self.patch_size					,num_channels=self.num_channels					,embed_dim=self.embed_dim					,depths=self.depths					,num_heads=self.num_heads					,window_size=self.window_size					,mlp_ratio=self.mlp_ratio					,qkv_bias=self.qkv_bias					,hidden_dropout_prob=self.hidden_dropout_prob					,attention_probs_dropout_prob=self.attention_probs_dropout_prob					,drop_path_rate=self.drop_path_rate					,hidden_act=self.hidden_act					,use_absolute_embeddings=self.use_absolute_embeddings					,path_norm=self.patch_norm					,layer_norm_eps=self.layer_norm_eps					,initializer_range=self.initializer_range					,encoder_stride=self.encoder_stride					,)






 def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase       )						->    Union[str, Any]:
     snake_case__      :str         =   SwinvaModel(config=lowerCamelCase_       )
     model.to(lowerCamelCase_       )
     model.eval()
     snake_case__      :int         =   model(lowerCamelCase_       )

     snake_case__      :List[Any]         =   ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths       ) - 1))
     snake_case__      :Dict         =   int(config.embed_dim * 2 ** (len(config.depths       ) - 1)       )

     self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, expected_seq_len, expected_dim)       )






 def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase       )						->    Union[str, Any]:
     snake_case__      :List[str]         =   SwinvaForMaskedImageModeling(config=lowerCamelCase_       )
     model.to(lowerCamelCase_       )
     model.eval()
     snake_case__      :Dict         =   model(lowerCamelCase_       )
     self.parent.assertEqual(
         result.logits.shape					,(self.batch_size, self.num_channels, self.image_size, self.image_size)       )

     # test greyscale images
     snake_case__      :Tuple         =   1
     snake_case__      :List[Any]         =   SwinvaForMaskedImageModeling(lowerCamelCase_       )
     model.to(lowerCamelCase_       )
     model.eval()

     snake_case__      :List[str]         =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
     snake_case__      :Tuple         =   model(lowerCamelCase_       )
     self.parent.assertEqual(result.logits.shape					,(self.batch_size, 1, self.image_size, self.image_size)       )






 def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase       )						->    Tuple:
     snake_case__      :Tuple         =   self.type_sequence_label_size
     snake_case__      :Optional[Any]         =   SwinvaForImageClassification(lowerCamelCase_       )
     model.to(lowerCamelCase_       )
     model.eval()
     snake_case__      :Optional[int]         =   model(lowerCamelCase_					,labels=lowerCamelCase_       )
     self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.type_sequence_label_size)       )






 def    lowerCAmelCase_							(		self       )						->    List[Any]:
     snake_case__      :List[str]         =   self.prepare_config_and_inputs()
     snake_case__      :Optional[int]         =   config_and_inputs
     snake_case__      :int         =   {'''pixel_values''': pixel_values}
     return config, inputs_dict




@require_torch
class 					_snake_case						(			_UpperCAmelCase       ,    _UpperCAmelCase       ,    unittest.TestCase ):
 _A          =   (
     (SwinvaModel, SwinvaForImageClassification, SwinvaForMaskedImageModeling) if is_torch_available() else ()
 )
 _A          =   (
     {'feature-extraction': SwinvaModel, 'image-classification': SwinvaForImageClassification}
     if is_torch_available()
     else {}
 )

 _A          =   False
 _A          =   False
 _A          =   False
 _A          =   False






 def    lowerCAmelCase_							(		self       )						->    Union[str, Any]:
     snake_case__      :int         =   SwinvaModelTester(self       )
     snake_case__      :Dict         =   ConfigTester(self					,config_class=lowerCamelCase_					,embed_dim=37       )






 def    lowerCAmelCase_							(		self       )						->    List[str]:
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()






 def    lowerCAmelCase_							(		self       )						->    List[str]:
     snake_case__      :Union[str, Any]         =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowerCamelCase_       )






 @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0."       )
 def    lowerCAmelCase_							(		self       )						->    List[Any]:
     pass






 @unittest.skip(reason="Swinv2 does not use inputs_embeds"       )
 def    lowerCAmelCase_							(		self       )						->    Any:
     pass






 def    lowerCAmelCase_							(		self       )						->    Union[str, Any]:
     snake_case__      :Optional[int]         =   self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         snake_case__      :List[Any]         =   model_class(lowerCamelCase_       )
         self.assertIsInstance(model.get_input_embeddings()					,(nn.Module)       )
         snake_case__      :str         =   model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(lowerCamelCase_					,nn.Linear       )       )






 def    lowerCAmelCase_							(		self       )						->    List[Any]:
     snake_case__      :Tuple         =   self.model_tester.prepare_config_and_inputs_for_common()

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

         snake_case__      :Union[str, Any]         =   ['''pixel_values''']
         self.assertListEqual(arg_names[:1]					,lowerCamelCase_       )






 def    lowerCAmelCase_							(		self       )						->    Optional[Any]:
     snake_case__      :str         =   self.model_tester.prepare_config_and_inputs_for_common()
     snake_case__      :Tuple         =   True

     for model_class in self.all_model_classes:
         snake_case__      :List[str]         =   True
         snake_case__      :List[Any]         =   False
         snake_case__      :Any         =   True
         snake_case__      :Any         =   model_class(lowerCamelCase_       )
         model.to(lowerCamelCase_       )
         model.eval()
         with torch.no_grad():
             snake_case__      :str         =   model(**self._prepare_for_class(lowerCamelCase_					,lowerCamelCase_       )       )
         snake_case__      :Tuple         =   outputs.attentions
         snake_case__      :Tuple         =   len(self.model_tester.depths       )
         self.assertEqual(len(lowerCamelCase_       )					,lowerCamelCase_       )

         # check that output_attentions also work using config
         del inputs_dict["output_attentions"]
         snake_case__      :Union[str, Any]         =   True
         snake_case__      :Optional[int]         =   config.window_size**2
         snake_case__      :str         =   model_class(lowerCamelCase_       )
         model.to(lowerCamelCase_       )
         model.eval()
         with torch.no_grad():
             snake_case__      :Union[str, Any]         =   model(**self._prepare_for_class(lowerCamelCase_					,lowerCamelCase_       )       )
         snake_case__      :Dict         =   outputs.attentions
         self.assertEqual(len(lowerCamelCase_       )					,lowerCamelCase_       )

         self.assertListEqual(
             list(attentions[0].shape[-3:]       )					,[self.model_tester.num_heads[0], window_size_squared, window_size_squared]					,)
         snake_case__      :Dict         =   len(lowerCamelCase_       )

         # Check attention is always last and order is fine
         snake_case__      :int         =   True
         snake_case__      :Dict         =   True
         snake_case__      :int         =   model_class(lowerCamelCase_       )
         model.to(lowerCamelCase_       )
         model.eval()
         with torch.no_grad():
             snake_case__      :List[Any]         =   model(**self._prepare_for_class(lowerCamelCase_					,lowerCamelCase_       )       )

         if hasattr(self.model_tester					,"num_hidden_states_types"       ):
             snake_case__      :Any         =   self.model_tester.num_hidden_states_types
         else:
             # also another +1 for reshaped_hidden_states
             snake_case__      :Optional[Any]         =   2
         self.assertEqual(out_len + added_hidden_states					,len(lowerCamelCase_       )       )

         snake_case__      :Tuple         =   outputs.attentions

         self.assertEqual(len(lowerCamelCase_       )					,lowerCamelCase_       )

         self.assertListEqual(
             list(self_attentions[0].shape[-3:]       )					,[self.model_tester.num_heads[0], window_size_squared, window_size_squared]					,)






 def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase       )						->    Optional[Any]:
     snake_case__      :Union[str, Any]         =   model_class(lowerCamelCase_       )
     model.to(lowerCamelCase_       )
     model.eval()

     with torch.no_grad():
         snake_case__      :Optional[int]         =   model(**self._prepare_for_class(lowerCamelCase_					,lowerCamelCase_       )       )

     snake_case__      :Dict         =   outputs.hidden_states

     snake_case__      :str         =   getattr(
         self.model_tester					,"expected_num_hidden_layers"					,len(self.model_tester.depths       ) + 1       )
     self.assertEqual(len(lowerCamelCase_       )					,lowerCamelCase_       )

     # Swinv2 has a different seq_length
     snake_case__      :List[str]         =   (
         config.patch_size
         if isinstance(config.patch_size					,collections.abc.Iterable       )
         else (config.patch_size, config.patch_size)
     )

     snake_case__      :Dict         =   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

     self.assertListEqual(
         list(hidden_states[0].shape[-2:]       )					,[num_patches, self.model_tester.embed_dim]					,)

     snake_case__      :Any         =   outputs.reshaped_hidden_states
     self.assertEqual(len(lowerCamelCase_       )					,lowerCamelCase_       )

     snake_case__      :Any         =   reshaped_hidden_states[0].shape
     snake_case__      :Optional[int]         =   (
         reshaped_hidden_states[0].view(lowerCamelCase_					,lowerCamelCase_					,height * width       ).permute(0					,2					,1       )
     )
     self.assertListEqual(
         list(reshaped_hidden_states.shape[-2:]       )					,[num_patches, self.model_tester.embed_dim]					,)






 def    lowerCAmelCase_							(		self       )						->    Optional[int]:
     snake_case__      :Union[str, Any]         =   self.model_tester.prepare_config_and_inputs_for_common()

     snake_case__      :List[str]         =   (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size					,collections.abc.Iterable       )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )

     for model_class in self.all_model_classes:
         snake_case__      :Tuple         =   True
         self.check_hidden_states_output(lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_       )

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         snake_case__      :List[str]         =   True

         self.check_hidden_states_output(lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_       )






 def    lowerCAmelCase_							(		self       )						->    Optional[int]:
     snake_case__      :str         =   self.model_tester.prepare_config_and_inputs_for_common()
     snake_case__      :List[Any]         =   3

     snake_case__      :int         =   (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size					,collections.abc.Iterable       )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )
     snake_case__      :Optional[Any]         =   (
         config.patch_size
         if isinstance(config.patch_size					,collections.abc.Iterable       )
         else (config.patch_size, config.patch_size)
     )

     snake_case__      :Optional[int]         =   image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
     snake_case__      :Optional[int]         =   image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

     for model_class in self.all_model_classes:
         snake_case__      :Optional[Any]         =   True
         self.check_hidden_states_output(lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,(padded_height, padded_width)       )

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         snake_case__      :Optional[Any]         =   True
         self.check_hidden_states_output(lowerCamelCase_					,lowerCamelCase_					,lowerCamelCase_					,(padded_height, padded_width)       )






 def    lowerCAmelCase_							(		self       )						->    List[Any]:
     snake_case__      :Tuple         =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase_       )






 def    lowerCAmelCase_							(		self       )						->    Optional[Any]:
     snake_case__      :Tuple         =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowerCamelCase_       )






 @slow
 def    lowerCAmelCase_							(		self       )						->    List[str]:




     for model_name in SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case__      :Tuple         =   SwinvaModel.from_pretrained(lowerCamelCase_       )
         self.assertIsNotNone(lowerCamelCase_       )






 def    lowerCAmelCase_							(		self       )						->    Any:
     snake_case__      :List[str]         =   self.model_tester.prepare_config_and_inputs_for_common()

     snake_case__      :Tuple         =   _config_zero_init(lowerCamelCase_       )
     for model_class in self.all_model_classes:
         snake_case__      :Dict         =   model_class(config=lowerCamelCase_       )
         for name, param in model.named_parameters():
             if "embeddings" not in name and "logit_scale" not in name and param.requires_grad:
                 self.assertIn(
                     ((param.data.mean() * 1E9).round() / 1E9).item()					,[0.0, 1.0]					,msg=f'Parameter {name} of model {model_class} seems not properly initialized'					,)




@require_vision
@require_torch
class 					_snake_case						(			unittest.TestCase ):






 @cached_property
 def    lowerCAmelCase_							(		self       )						->    List[Any]:
     return (
         AutoImageProcessor.from_pretrained("microsoft/swinv2-tiny-patch4-window8-256"       )
         if is_vision_available()
         else None
     )






 @slow
 def    lowerCAmelCase_							(		self       )						->    Tuple:
     snake_case__      :Optional[int]         =   SwinvaForImageClassification.from_pretrained("microsoft/swinv2-tiny-patch4-window8-256"       ).to(
         lowerCamelCase_       )
     snake_case__      :Dict         =   self.default_image_processor

     snake_case__      :Optional[Any]         =   Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
     snake_case__      :List[str]         =   image_processor(images=lowerCamelCase_					,return_tensors="pt"       ).to(lowerCamelCase_       )

     # forward pass
     with torch.no_grad():
         snake_case__      :Tuple         =   model(**lowerCamelCase_       )

     # verify the logits
     snake_case__      :List[str]         =   torch.Size((1, 1_000)       )
     self.assertEqual(outputs.logits.shape					,lowerCamelCase_       )
     snake_case__      :int         =   torch.tensor([-0.3947, -0.4306, 0.0026]       ).to(lowerCamelCase_       )
     self.assertTrue(torch.allclose(outputs.logits[0, :3]					,lowerCamelCase_					,atol=1E-4       )       )







import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, 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 					_snake_case						(			unittest.TestCase ):






	def    lowerCAmelCase_							(		self       )						->    int:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()






	def    lowerCAmelCase_							(		self       )						->    str:
					snake_case__    ,						snake_case__      :Tuple         =   FlaxControlNetModel.from_pretrained(
					    "lllyasviel/sd-controlnet-canny"					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__    ,						snake_case__      :Any         =   FlaxStableDiffusionControlNetPipeline.from_pretrained(
					    "runwayml/stable-diffusion-v1-5"					,controlnet=UpperCamelCase					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__      :List[str]         =   controlnet_params

					snake_case__      :Union[str, Any]         =   "bird"
					snake_case__      :Optional[int]         =   jax.device_count()
					snake_case__      :Tuple         =   pipe.prepare_text_inputs([prompts] * num_samples       )

					snake_case__      :Union[str, Any]         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"       )
					snake_case__      :str         =   pipe.prepare_image_inputs([canny_image] * num_samples       )

					snake_case__      :List[str]         =   jax.random.PRNGKey(0       )
					snake_case__      :str         =   jax.random.split(UpperCamelCase					,jax.device_count()       )

					snake_case__      :int         =   replicate(UpperCamelCase       )
					snake_case__      :Any         =   shard(UpperCamelCase       )
					snake_case__      :Any         =   shard(UpperCamelCase       )

					snake_case__      :str         =   pipe(
					    prompt_ids=UpperCamelCase					,image=UpperCamelCase					,params=UpperCamelCase					,prng_seed=UpperCamelCase					,num_inference_steps=50					,jit=UpperCamelCase					,).images
					assert images.shape == (jax.device_count(), 1, 768, 512, 3)

					snake_case__      :List[str]         =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
					snake_case__      :Any         =   images[0, 253:256, 253:256, -1]

					snake_case__      :Union[str, Any]         =   jnp.asarray(jax.device_get(image_slice.flatten()       )       )
					snake_case__      :List[Any]         =   jnp.array(
					    [0.167969, 0.116699, 0.081543, 0.154297, 0.132812, 0.108887, 0.169922, 0.169922, 0.205078]       )
					print(f'output_slice: {output_slice}'       )
					assert jnp.abs(output_slice - expected_slice       ).max() < 1E-2






	def    lowerCAmelCase_							(		self       )						->    Optional[int]:
					snake_case__    ,						snake_case__      :List[str]         =   FlaxControlNetModel.from_pretrained(
					    "lllyasviel/sd-controlnet-openpose"					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__    ,						snake_case__      :Optional[Any]         =   FlaxStableDiffusionControlNetPipeline.from_pretrained(
					    "runwayml/stable-diffusion-v1-5"					,controlnet=UpperCamelCase					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__      :str         =   controlnet_params

					snake_case__      :int         =   "Chef in the kitchen"
					snake_case__      :List[Any]         =   jax.device_count()
					snake_case__      :Dict         =   pipe.prepare_text_inputs([prompts] * num_samples       )

					snake_case__      :Any         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png"       )
					snake_case__      :Optional[int]         =   pipe.prepare_image_inputs([pose_image] * num_samples       )

					snake_case__      :List[str]         =   jax.random.PRNGKey(0       )
					snake_case__      :Any         =   jax.random.split(UpperCamelCase					,jax.device_count()       )

					snake_case__      :Dict         =   replicate(UpperCamelCase       )
					snake_case__      :Tuple         =   shard(UpperCamelCase       )
					snake_case__      :Optional[int]         =   shard(UpperCamelCase       )

					snake_case__      :Optional[Any]         =   pipe(
					    prompt_ids=UpperCamelCase					,image=UpperCamelCase					,params=UpperCamelCase					,prng_seed=UpperCamelCase					,num_inference_steps=50					,jit=UpperCamelCase					,).images
					assert images.shape == (jax.device_count(), 1, 768, 512, 3)

					snake_case__      :int         =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
					snake_case__      :List[str]         =   images[0, 253:256, 253:256, -1]

					snake_case__      :Tuple         =   jnp.asarray(jax.device_get(image_slice.flatten()       )       )
					snake_case__      :List[str]         =   jnp.array(
					    [[0.271484, 0.261719, 0.275391, 0.277344, 0.279297, 0.291016, 0.294922, 0.302734, 0.302734]]       )
					print(f'output_slice: {output_slice}'       )
					assert jnp.abs(output_slice - expected_slice       ).max() < 1E-2