"""simple docstring"""






import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)






def 					snake_case__       (		__lowerCamelCase    :   Optional[Any]     ):
			"""simple docstring"""
			lowerCamelCase__		:		int		     =VideoMAEConfig()

			set_architecture_configs(__lowerCamelCase							,		__lowerCamelCase     )

			if "finetuned" not in model_name:
						lowerCamelCase__		:		Optional[int]		     =False

			if "finetuned" in model_name:
						lowerCamelCase__		:		int		     ='''huggingface/label-files'''
						if "kinetics" in model_name:
									lowerCamelCase__		:		str		     =400
									lowerCamelCase__		:		Tuple		     ='''kinetics400-id2label.json'''
						elif "ssv2" in model_name:
									lowerCamelCase__		:		Any		     =174
									lowerCamelCase__		:		str		     ='''something-something-v2-id2label.json'''
						else:
									raise ValueError('''Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'''     )
						lowerCamelCase__		:		int		     =json.load(open(hf_hub_download(__lowerCamelCase							,		__lowerCamelCase							,		repo_type='''dataset'''     )							,		'''r'''     )     )
						lowerCamelCase__		:		Any		     ={int(__lowerCamelCase     ): v for k, v in idalabel.items()}
						lowerCamelCase__		:		Union[str, Any]		     =idalabel
						lowerCamelCase__		:		Optional[int]		     ={v: k for k, v in idalabel.items()}

			return config






def 					snake_case__       (		__lowerCamelCase    :   int							,		__lowerCamelCase    :   Dict     ):
			"""simple docstring"""
			if "small" in model_name:
						lowerCamelCase__		:		Any		     =384
						lowerCamelCase__		:		Optional[Any]		     =1536
						lowerCamelCase__		:		List[str]		     =12
						lowerCamelCase__		:		List[Any]		     =16
						lowerCamelCase__		:		int		     =12
						lowerCamelCase__		:		List[Any]		     =3
						lowerCamelCase__		:		Union[str, Any]		     =192
						lowerCamelCase__		:		Optional[int]		     =768
			elif "large" in model_name:
						lowerCamelCase__		:		Optional[int]		     =1024
						lowerCamelCase__		:		List[str]		     =4096
						lowerCamelCase__		:		Optional[Any]		     =24
						lowerCamelCase__		:		Any		     =16
						lowerCamelCase__		:		Optional[int]		     =12
						lowerCamelCase__		:		Tuple		     =8
						lowerCamelCase__		:		List[str]		     =512
						lowerCamelCase__		:		Any		     =2048
			elif "huge" in model_name:
						lowerCamelCase__		:		Dict		     =1280
						lowerCamelCase__		:		str		     =5120
						lowerCamelCase__		:		List[str]		     =32
						lowerCamelCase__		:		str		     =16
						lowerCamelCase__		:		Tuple		     =12
						lowerCamelCase__		:		Optional[Any]		     =8
						lowerCamelCase__		:		Tuple		     =640
						lowerCamelCase__		:		Dict		     =2560
			elif "base" not in model_name:
						raise ValueError('''Model name should include either "small", "base", "large", or "huge"'''     )






def 					snake_case__       (		__lowerCamelCase    :   Optional[int]     ):
			"""simple docstring"""
			if "encoder." in name:
						lowerCamelCase__		:		List[Any]		     =name.replace('''encoder.'''							,		''''''     )
			if "cls_token" in name:
						lowerCamelCase__		:		Any		     =name.replace('''cls_token'''							,		'''videomae.embeddings.cls_token'''     )
			if "decoder_pos_embed" in name:
						lowerCamelCase__		:		Dict		     =name.replace('''decoder_pos_embed'''							,		'''decoder.decoder_pos_embed'''     )
			if "pos_embed" in name and "decoder" not in name:
						lowerCamelCase__		:		List[Any]		     =name.replace('''pos_embed'''							,		'''videomae.embeddings.position_embeddings'''     )
			if "patch_embed.proj" in name:
						lowerCamelCase__		:		int		     =name.replace('''patch_embed.proj'''							,		'''videomae.embeddings.patch_embeddings.projection'''     )
			if "patch_embed.norm" in name:
						lowerCamelCase__		:		List[str]		     =name.replace('''patch_embed.norm'''							,		'''videomae.embeddings.norm'''     )
			if "decoder.blocks" in name:
						lowerCamelCase__		:		Optional[Any]		     =name.replace('''decoder.blocks'''							,		'''decoder.decoder_layers'''     )
			if "blocks" in name:
						lowerCamelCase__		:		List[Any]		     =name.replace('''blocks'''							,		'''videomae.encoder.layer'''     )
			if "attn.proj" in name:
						lowerCamelCase__		:		Tuple		     =name.replace('''attn.proj'''							,		'''attention.output.dense'''     )
			if "attn" in name and "bias" not in name:
						lowerCamelCase__		:		List[Any]		     =name.replace('''attn'''							,		'''attention.self'''     )
			if "attn" in name:
						lowerCamelCase__		:		List[str]		     =name.replace('''attn'''							,		'''attention.attention'''     )
			if "norm1" in name:
						lowerCamelCase__		:		int		     =name.replace('''norm1'''							,		'''layernorm_before'''     )
			if "norm2" in name:
						lowerCamelCase__		:		List[str]		     =name.replace('''norm2'''							,		'''layernorm_after'''     )
			if "mlp.fc1" in name:
						lowerCamelCase__		:		str		     =name.replace('''mlp.fc1'''							,		'''intermediate.dense'''     )
			if "mlp.fc2" in name:
						lowerCamelCase__		:		Optional[int]		     =name.replace('''mlp.fc2'''							,		'''output.dense'''     )
			if "decoder_embed" in name:
						lowerCamelCase__		:		Any		     =name.replace('''decoder_embed'''							,		'''decoder.decoder_embed'''     )
			if "decoder_norm" in name:
						lowerCamelCase__		:		Union[str, Any]		     =name.replace('''decoder_norm'''							,		'''decoder.decoder_norm'''     )
			if "decoder_pred" in name:
						lowerCamelCase__		:		List[str]		     =name.replace('''decoder_pred'''							,		'''decoder.decoder_pred'''     )
			if "norm.weight" in name and "decoder" not in name and "fc" not in name:
						lowerCamelCase__		:		Union[str, Any]		     =name.replace('''norm.weight'''							,		'''videomae.layernorm.weight'''     )
			if "norm.bias" in name and "decoder" not in name and "fc" not in name:
						lowerCamelCase__		:		str		     =name.replace('''norm.bias'''							,		'''videomae.layernorm.bias'''     )
			if "head" in name and "decoder" not in name:
						lowerCamelCase__		:		Any		     =name.replace('''head'''							,		'''classifier'''     )

			return name






def 					snake_case__       (		__lowerCamelCase    :   Tuple							,		__lowerCamelCase    :   Any     ):
			"""simple docstring"""
			for key in orig_state_dict.copy().keys():
						lowerCamelCase__		:		Tuple		     =orig_state_dict.pop(__lowerCamelCase     )

						if key.startswith('''encoder.'''     ):
									lowerCamelCase__		:		int		     =key.replace('''encoder.'''							,		''''''     )

						if "qkv" in key:
									lowerCamelCase__		:		Tuple		     =key.split('''.'''     )
									if key.startswith('''decoder.blocks'''     ):
												lowerCamelCase__		:		str		     =config.decoder_hidden_size
												lowerCamelCase__		:		List[str]		     =int(key_split[2]     )
												lowerCamelCase__		:		Tuple		     ='''decoder.decoder_layers.'''
												if "weight" in key:
															lowerCamelCase__		:		Optional[Any]		     =val[:dim, :]
															lowerCamelCase__		:		List[str]		     =val[dim : dim * 2, :]
															lowerCamelCase__		:		int		     =val[-dim:, :]
									else:
												lowerCamelCase__		:		List[str]		     =config.hidden_size
												lowerCamelCase__		:		Union[str, Any]		     =int(key_split[1]     )
												lowerCamelCase__		:		int		     ='''videomae.encoder.layer.'''
												if "weight" in key:
															lowerCamelCase__		:		int		     =val[:dim, :]
															lowerCamelCase__		:		Tuple		     =val[dim : dim * 2, :]
															lowerCamelCase__		:		str		     =val[-dim:, :]
						else:
									lowerCamelCase__		:		str		     =val

			return orig_state_dict






def 					snake_case__       (		):
			"""simple docstring"""
			lowerCamelCase__		:		str		     =hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''							,		filename='''eating_spaghetti.npy'''							,		repo_type='''dataset'''     )
			lowerCamelCase__		:		List[str]		     =np.load(__lowerCamelCase     )
			return list(__lowerCamelCase     )






def 					snake_case__       (		__lowerCamelCase    :   Any							,		__lowerCamelCase    :   Optional[Any]							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   Any     ):
			"""simple docstring"""
			lowerCamelCase__		:		List[str]		     =get_videomae_config(__lowerCamelCase     )

			if "finetuned" in model_name:
						lowerCamelCase__		:		List[str]		     =VideoMAEForVideoClassification(__lowerCamelCase     )
			else:
						lowerCamelCase__		:		str		     =VideoMAEForPreTraining(__lowerCamelCase     )

			# download original checkpoint, hosted on Google Drive
			lowerCamelCase__		:		List[str]		     ='''pytorch_model.bin'''
			gdown.cached_download(__lowerCamelCase							,		__lowerCamelCase							,		quiet=__lowerCamelCase     )
			lowerCamelCase__		:		Tuple		     =torch.load(__lowerCamelCase							,		map_location='''cpu'''     )
			if "model" in files:
						lowerCamelCase__		:		List[Any]		     =files['''model''']
			else:
						lowerCamelCase__		:		Dict		     =files['''module''']
			lowerCamelCase__		:		Tuple		     =convert_state_dict(__lowerCamelCase							,		__lowerCamelCase     )

			model.load_state_dict(__lowerCamelCase     )
			model.eval()

			# verify model on basic input
			lowerCamelCase__		:		Optional[int]		     =VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,		image_std=[0.5, 0.5, 0.5]     )
			lowerCamelCase__		:		str		     =prepare_video()
			lowerCamelCase__		:		Optional[int]		     =image_processor(__lowerCamelCase							,		return_tensors='''pt'''     )

			if "finetuned" not in model_name:
						lowerCamelCase__		:		Tuple		     =hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''							,		filename='''bool_masked_pos.pt'''     )
						lowerCamelCase__		:		List[Any]		     =torch.load(__lowerCamelCase     )

			lowerCamelCase__		:		str		     =model(**__lowerCamelCase     )
			lowerCamelCase__		:		List[str]		     =outputs.logits

			lowerCamelCase__		:		Optional[int]		     =[
			    '''videomae-small-finetuned-kinetics''',
			    '''videomae-small-finetuned-ssv2''',
			    # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
			    '''videomae-base-short''',
			    '''videomae-base-short-finetuned-kinetics''',
			    '''videomae-base''',
			    '''videomae-base-finetuned-kinetics''',
			    '''videomae-large''',
			    '''videomae-large-finetuned-kinetics''',
			    '''videomae-huge-finetuned-kinetics''',
			    # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
			    '''videomae-base-short-ssv2''',
			    '''videomae-base-short-finetuned-ssv2''',
			    '''videomae-base-ssv2''',
			    '''videomae-base-finetuned-ssv2''',
			]

			# NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
			if model_name == "videomae-small-finetuned-kinetics":
						lowerCamelCase__		:		int		     =torch.Size([1, 400]     )
						lowerCamelCase__		:		Dict		     =torch.tensor([-0.92_91, -0.40_61, -0.93_07]     )
			elif model_name == "videomae-small-finetuned-ssv2":
						lowerCamelCase__		:		Optional[Any]		     =torch.Size([1, 174]     )
						lowerCamelCase__		:		str		     =torch.tensor([0.26_71, -0.46_89, -0.82_35]     )
			elif model_name == "videomae-base":
						lowerCamelCase__		:		Dict		     =torch.Size([1, 1408, 1536]     )
						lowerCamelCase__		:		Tuple		     =torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]     )
			elif model_name == "videomae-base-short":
						lowerCamelCase__		:		List[str]		     =torch.Size([1, 1408, 1536]     )
						lowerCamelCase__		:		List[str]		     =torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]     )
						# we verified the loss both for normalized and unnormalized targets for this one
						lowerCamelCase__		:		Any		     =torch.tensor([0.51_42]     ) if config.norm_pix_loss else torch.tensor([0.64_69]     )
			elif model_name == "videomae-large":
						lowerCamelCase__		:		Optional[Any]		     =torch.Size([1, 1408, 1536]     )
						lowerCamelCase__		:		Optional[Any]		     =torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]     )
			elif model_name == "videomae-large-finetuned-kinetics":
						lowerCamelCase__		:		List[str]		     =torch.Size([1, 400]     )
						lowerCamelCase__		:		Tuple		     =torch.tensor([0.07_71, 0.00_11, -0.36_25]     )
			elif model_name == "videomae-huge-finetuned-kinetics":
						lowerCamelCase__		:		List[Any]		     =torch.Size([1, 400]     )
						lowerCamelCase__		:		List[Any]		     =torch.tensor([0.24_33, 0.16_32, -0.48_94]     )
			elif model_name == "videomae-base-short-finetuned-kinetics":
						lowerCamelCase__		:		Any		     =torch.Size([1, 400]     )
						lowerCamelCase__		:		List[Any]		     =torch.tensor([0.65_88, 0.09_90, -0.24_93]     )
			elif model_name == "videomae-base-finetuned-kinetics":
						lowerCamelCase__		:		List[Any]		     =torch.Size([1, 400]     )
						lowerCamelCase__		:		List[str]		     =torch.tensor([0.36_69, -0.06_88, -0.24_21]     )
			elif model_name == "videomae-base-short-ssv2":
						lowerCamelCase__		:		Optional[int]		     =torch.Size([1, 1408, 1536]     )
						lowerCamelCase__		:		Optional[int]		     =torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]     )
			elif model_name == "videomae-base-short-finetuned-ssv2":
						lowerCamelCase__		:		Dict		     =torch.Size([1, 174]     )
						lowerCamelCase__		:		Tuple		     =torch.tensor([-0.05_37, -0.15_39, -0.32_66]     )
			elif model_name == "videomae-base-ssv2":
						lowerCamelCase__		:		int		     =torch.Size([1, 1408, 1536]     )
						lowerCamelCase__		:		Union[str, Any]		     =torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]     )
			elif model_name == "videomae-base-finetuned-ssv2":
						lowerCamelCase__		:		int		     =torch.Size([1, 174]     )
						lowerCamelCase__		:		List[Any]		     =torch.tensor([0.19_61, -0.83_37, -0.63_89]     )
			else:
						raise ValueError(f'''Model name not supported. Should be one of {model_names}'''     )

			# verify logits
			assert logits.shape == expected_shape
			if "finetuned" in model_name:
						assert torch.allclose(logits[0, :3]							,		__lowerCamelCase							,		atol=1e-4     )
			else:
						print('''Logits:'''							,		logits[0, :3, :3]     )
						assert torch.allclose(logits[0, :3, :3]							,		__lowerCamelCase							,		atol=1e-4     )
			print('''Logits ok!'''     )

			# verify loss, if applicable
			if model_name == "videomae-base-short":
						lowerCamelCase__		:		List[str]		     =outputs.loss
						assert torch.allclose(__lowerCamelCase							,		__lowerCamelCase							,		atol=1e-4     )
						print('''Loss ok!'''     )

			if pytorch_dump_folder_path is not None:
						print(f'''Saving model and image processor to {pytorch_dump_folder_path}'''     )
						image_processor.save_pretrained(__lowerCamelCase     )
						model.save_pretrained(__lowerCamelCase     )

			if push_to_hub:
						print('''Pushing to the hub...'''     )
						model.push_to_hub(__lowerCamelCase							,		organization='''nielsr'''     )


if __name__ == "__main__":
		_lowercase			:			int										=						argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--checkpoint_url",
		    default="https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4",
		    type=str,
		    help=(
		        "URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"
		        " download link."
		    ),
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path",
		    default="/Users/nielsrogge/Documents/VideoMAE/Test",
		    type=str,
		    help="Path to the output PyTorch model directory.",
		)
		parser.add_argument("--model_name", default="videomae-base", type=str, help="Name of the model.")
		parser.add_argument(
		    "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
		)

		_lowercase			:			List[str]										=						parser.parse_args()
		convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)









"""simple docstring"""






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


					'''simple docstring'''





					def 		snake_case	(	self							:  Union[str, Any]					)-> Dict:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()



					def 		snake_case	(	self							:  str					)-> Any:
								lowerCamelCase__     ,       lowerCamelCase__		:		Optional[Any]		     =FlaxControlNetModel.from_pretrained(
								    '''lllyasviel/sd-controlnet-canny''',  from_pt=lowerCamelCase,  dtype=jnp.bfloataa					)
								lowerCamelCase__     ,       lowerCamelCase__		:		Optional[int]		     =FlaxStableDiffusionControlNetPipeline.from_pretrained(
								    '''runwayml/stable-diffusion-v1-5''',  controlnet=lowerCamelCase,  from_pt=lowerCamelCase,  dtype=jnp.bfloataa					)
								lowerCamelCase__		:		Optional[int]		     =controlnet_params

								lowerCamelCase__		:		Dict		     ='''bird'''
								lowerCamelCase__		:		List[str]		     =jax.device_count()
								lowerCamelCase__		:		Optional[Any]		     =pipe.prepare_text_inputs([prompts] * num_samples					)

								lowerCamelCase__		:		Dict		     =load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''					)
								lowerCamelCase__		:		List[Any]		     =pipe.prepare_image_inputs([canny_image] * num_samples					)

								lowerCamelCase__		:		Optional[int]		     =jax.random.PRNGKey(0					)
								lowerCamelCase__		:		Dict		     =jax.random.split(lowerCamelCase,  jax.device_count()					)

								lowerCamelCase__		:		Tuple		     =replicate(lowerCamelCase					)
								lowerCamelCase__		:		Tuple		     =shard(lowerCamelCase					)
								lowerCamelCase__		:		Optional[int]		     =shard(lowerCamelCase					)

								lowerCamelCase__		:		Tuple		     =pipe(
								    prompt_ids=lowerCamelCase,  image=lowerCamelCase,  params=lowerCamelCase,  prng_seed=lowerCamelCase,  num_inference_steps=50,  jit=lowerCamelCase,  ).images
								assert images.shape == (jax.device_count(), 1, 768, 512, 3)

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

								lowerCamelCase__		:		Optional[Any]		     =jnp.asarray(jax.device_get(image_slice.flatten()					)					)
								lowerCamelCase__		:		Dict		     =jnp.array(
								    [0.167_969, 0.116_699, 0.081_543, 0.154_297, 0.132_812, 0.108_887, 0.169_922, 0.169_922, 0.205_078]					)
								print(F'''output_slice: {output_slice}'''					)
								assert jnp.abs(output_slice - expected_slice					).max() < 1E-2




					def 		snake_case	(	self							:  Optional[int]					)-> Optional[int]:
								lowerCamelCase__     ,       lowerCamelCase__		:		Dict		     =FlaxControlNetModel.from_pretrained(
								    '''lllyasviel/sd-controlnet-openpose''',  from_pt=lowerCamelCase,  dtype=jnp.bfloataa					)
								lowerCamelCase__     ,       lowerCamelCase__		:		List[Any]		     =FlaxStableDiffusionControlNetPipeline.from_pretrained(
								    '''runwayml/stable-diffusion-v1-5''',  controlnet=lowerCamelCase,  from_pt=lowerCamelCase,  dtype=jnp.bfloataa					)
								lowerCamelCase__		:		Optional[Any]		     =controlnet_params

								lowerCamelCase__		:		int		     ='''Chef in the kitchen'''
								lowerCamelCase__		:		Optional[Any]		     =jax.device_count()
								lowerCamelCase__		:		Any		     =pipe.prepare_text_inputs([prompts] * num_samples					)

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

								lowerCamelCase__		:		Tuple		     =jax.random.PRNGKey(0					)
								lowerCamelCase__		:		Optional[Any]		     =jax.random.split(lowerCamelCase,  jax.device_count()					)

								lowerCamelCase__		:		int		     =replicate(lowerCamelCase					)
								lowerCamelCase__		:		List[Any]		     =shard(lowerCamelCase					)
								lowerCamelCase__		:		int		     =shard(lowerCamelCase					)

								lowerCamelCase__		:		Tuple		     =pipe(
								    prompt_ids=lowerCamelCase,  image=lowerCamelCase,  params=lowerCamelCase,  prng_seed=lowerCamelCase,  num_inference_steps=50,  jit=lowerCamelCase,  ).images
								assert images.shape == (jax.device_count(), 1, 768, 512, 3)

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

								lowerCamelCase__		:		int		     =jnp.asarray(jax.device_get(image_slice.flatten()					)					)
								lowerCamelCase__		:		Any		     =jnp.array(
								    [[0.271_484, 0.261_719, 0.275_391, 0.277_344, 0.279_297, 0.291_016, 0.294_922, 0.302_734, 0.302_734]]					)
								print(F'''output_slice: {output_slice}'''					)
								assert jnp.abs(output_slice - expected_slice					).max() < 1E-2




