import argparse

import torch

from transformers import (
    SpeechTaConfig,
    SpeechTaFeatureExtractor,
    SpeechTaForSpeechToSpeech,
    SpeechTaForSpeechToText,
    SpeechTaForTextToSpeech,
    SpeechTaProcessor,
    SpeechTaTokenizer,
    logging,
)
from transformers.tokenization_utils import AddedToken


logging.set_verbosity_info()
UpperCAmelCase    =					logging.get_logger("""transformers.models.speecht5""")

UpperCAmelCase    =					{
    '''speech_encoder_prenet.layer_norm''': '''speecht5.encoder.prenet.feature_projection.layer_norm''',
    '''speech_encoder_prenet.post_extract_proj''': '''speecht5.encoder.prenet.feature_projection.projection''',
    '''speech_encoder_prenet.pos_conv.0''': '''speecht5.encoder.prenet.pos_conv_embed.conv''',
    '''speech_encoder_prenet.mask_emb''': '''speecht5.encoder.prenet.masked_spec_embed''',
}
UpperCAmelCase    =					{
    '''text_encoder_prenet.encoder_prenet.0''': '''speecht5.encoder.prenet.embed_tokens''',
    '''text_encoder_prenet.encoder_prenet.1.alpha''': '''speecht5.encoder.prenet.encode_positions.alpha''',
}
UpperCAmelCase    =					{
    '''speech_decoder_prenet.decoder_prenet.0.0.prenet.0.0''': '''speecht5.decoder.prenet.layers.0''',
    '''speech_decoder_prenet.decoder_prenet.0.0.prenet.1.0''': '''speecht5.decoder.prenet.layers.1''',
    '''speech_decoder_prenet.decoder_prenet.0.1''': '''speecht5.decoder.prenet.final_layer''',
    '''speech_decoder_prenet.decoder_prenet.1.alpha''': '''speecht5.decoder.prenet.encode_positions.alpha''',
    '''speech_decoder_prenet.spkembs_layer.0''': '''speecht5.decoder.prenet.speaker_embeds_layer''',
}
UpperCAmelCase    =					{
    '''speech_decoder_postnet.feat_out''': '''speech_decoder_postnet.feat_out''',
    '''speech_decoder_postnet.prob_out''': '''speech_decoder_postnet.prob_out''',
    '''speech_decoder_postnet.postnet.postnet.0.0''': '''speech_decoder_postnet.layers.0.conv''',
    '''speech_decoder_postnet.postnet.postnet.0.1''': '''speech_decoder_postnet.layers.0.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.1.0''': '''speech_decoder_postnet.layers.1.conv''',
    '''speech_decoder_postnet.postnet.postnet.1.1''': '''speech_decoder_postnet.layers.1.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.2.0''': '''speech_decoder_postnet.layers.2.conv''',
    '''speech_decoder_postnet.postnet.postnet.2.1''': '''speech_decoder_postnet.layers.2.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.3.0''': '''speech_decoder_postnet.layers.3.conv''',
    '''speech_decoder_postnet.postnet.postnet.3.1''': '''speech_decoder_postnet.layers.3.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.4.0''': '''speech_decoder_postnet.layers.4.conv''',
    '''speech_decoder_postnet.postnet.postnet.4.1''': '''speech_decoder_postnet.layers.4.batch_norm''',
}
UpperCAmelCase    =					{
    '''text_decoder_prenet.embed_tokens''': '''speecht5.decoder.prenet.embed_tokens''',
}
UpperCAmelCase    =					{
    '''text_decoder_postnet.output_projection''': '''text_decoder_postnet.lm_head''',
}
UpperCAmelCase    =					{
    '''encoder.layers.*.self_attn.k_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.k_proj''',
    '''encoder.layers.*.self_attn.v_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.v_proj''',
    '''encoder.layers.*.self_attn.q_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.q_proj''',
    '''encoder.layers.*.self_attn.out_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.out_proj''',
    '''encoder.layers.*.self_attn_layer_norm''': '''speecht5.encoder.wrapped_encoder.layers.*.layer_norm''',
    '''encoder.layers.*.fc1''': '''speecht5.encoder.wrapped_encoder.layers.*.feed_forward.intermediate_dense''',
    '''encoder.layers.*.fc2''': '''speecht5.encoder.wrapped_encoder.layers.*.feed_forward.output_dense''',
    '''encoder.layers.*.final_layer_norm''': '''speecht5.encoder.wrapped_encoder.layers.*.final_layer_norm''',
    '''encoder.layer_norm''': '''speecht5.encoder.wrapped_encoder.layer_norm''',
    '''encoder.pos_emb.pe_k''': '''speecht5.encoder.wrapped_encoder.embed_positions.pe_k''',
}
UpperCAmelCase    =					{
    '''decoder.layers.*.self_attn.k_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.k_proj''',
    '''decoder.layers.*.self_attn.v_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.v_proj''',
    '''decoder.layers.*.self_attn.q_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.q_proj''',
    '''decoder.layers.*.self_attn.out_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.out_proj''',
    '''decoder.layers.*.self_attn_layer_norm''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn_layer_norm''',
    '''decoder.layers.*.encoder_attn.k_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.k_proj''',
    '''decoder.layers.*.encoder_attn.v_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.v_proj''',
    '''decoder.layers.*.encoder_attn.q_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.q_proj''',
    '''decoder.layers.*.encoder_attn.out_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.out_proj''',
    '''decoder.layers.*.encoder_attn_layer_norm''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn_layer_norm''',
    '''decoder.layers.*.fc1''': '''speecht5.decoder.wrapped_decoder.layers.*.feed_forward.intermediate_dense''',
    '''decoder.layers.*.fc2''': '''speecht5.decoder.wrapped_decoder.layers.*.feed_forward.output_dense''',
    '''decoder.layers.*.final_layer_norm''': '''speecht5.decoder.wrapped_decoder.layers.*.final_layer_norm''',
}
UpperCAmelCase    =					{
    **MAPPING_SPEECH_ENCODER_PRENET,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
    **MAPPING_TEXT_DECODER_PRENET,
    **MAPPING_TEXT_DECODER_POSTNET,
}
UpperCAmelCase    =					{
    **MAPPING_TEXT_ENCODER_PRENET,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
    **MAPPING_SPEECH_DECODER_PRENET,
    **MAPPING_SPEECH_DECODER_POSTNET,
}
UpperCAmelCase    =					{
    **MAPPING_SPEECH_ENCODER_PRENET,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
    **MAPPING_SPEECH_DECODER_PRENET,
    **MAPPING_SPEECH_DECODER_POSTNET,
}
UpperCAmelCase    =					[]
UpperCAmelCase    =					[
    '''encoder.version''',
    '''encoder.layers.*.norm_k.weight''',
    '''encoder.layers.*.norm_k.bias''',
    '''decoder.version''',
    '''decoder.layers.*.norm_k.weight''',
    '''decoder.layers.*.norm_k.bias''',
    '''decoder.pos_emb.pe_k''',
    '''speech_encoder_prenet.embed_positions._float_tensor''',
    '''text_decoder_prenet.embed_positions._float_tensor''',
]
UpperCAmelCase    =					IGNORE_KEYS + [
    '''encoder.proj''',
    '''text_encoder_prenet.*''',
    '''speech_decoder_prenet.*''',
    '''speech_decoder_postnet.*''',
]
UpperCAmelCase    =					IGNORE_KEYS + [
    '''encoder.proj''',
    '''speech_encoder_prenet.*''',
    '''text_decoder_prenet.*''',
    '''text_decoder_postnet.*''',
]
UpperCAmelCase    =					IGNORE_KEYS + [
    '''encoder.proj''',
    '''text_encoder_prenet.*''',
    '''text_decoder_prenet.*''',
    '''text_decoder_postnet.*''',
]






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						Tuple:




      """simple docstring"""

      for attribute in key.split('''.'''      ):
            snake_case_ =   getattr(lowerCAmelCase__      ,      lowerCAmelCase__      )

      if weight_type is not None:
            snake_case_ =   getattr(lowerCAmelCase__      ,      lowerCAmelCase__      ).shape
      else:
            snake_case_ =   hf_pointer.shape

      if hf_shape != value.shape:
            raise ValueError(
                f'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
                f''' {value.shape} for {full_name}'''      )

      if weight_type == "weight":
            snake_case_ =   value
      elif weight_type == "weight_g":
            snake_case_ =   value
      elif weight_type == "weight_v":
            snake_case_ =   value
      elif weight_type == "bias":
            snake_case_ =   value
      elif weight_type == "running_mean":
            snake_case_ =   value
      elif weight_type == "running_var":
            snake_case_ =   value
      elif weight_type == "num_batches_tracked":
            snake_case_ =   value
      else:
            snake_case_ =   value

      logger.info(f'''{key + ("." + weight_type if weight_type is not None else "")} was initialized from {full_name}.'''      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						Any:




      """simple docstring"""

      for key in ignore_keys:
            if key.endswith('''.*'''      ):
                  if name.startswith(key[:-1]      ):
                        return True
            elif ".*." in key:
                  snake_case_ ,					snake_case_ =   key.split('''.*.'''      )
                  if prefix in name and suffix in name:
                        return True
            elif key in name:
                  return True
      return False






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						int:




      """simple docstring"""

      snake_case_ =   []

      if task == "s2t":
            snake_case_ =   hf_model.speechta.encoder.prenet.feature_encoder
            snake_case_ =   MAPPING_S2T
            snake_case_ =   IGNORE_KEYS_S2T
      elif task == "t2s":
            snake_case_ =   None
            snake_case_ =   MAPPING_T2S
            snake_case_ =   IGNORE_KEYS_T2S
      elif task == "s2s":
            snake_case_ =   hf_model.speechta.encoder.prenet.feature_encoder
            snake_case_ =   MAPPING_S2S
            snake_case_ =   IGNORE_KEYS_S2S
      else:
            raise ValueError(f'''Unsupported task: {task}'''      )

      for name, value in fairseq_dict.items():
            if should_ignore(lowerCAmelCase__      ,      lowerCAmelCase__      ):
                  logger.info(f'''{name} was ignored'''      )
                  continue

            snake_case_ =   False
            if "conv_layers" in name:
                  load_conv_layer(
                      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      hf_model.config.feat_extract_norm == '''group'''      ,      )
                  snake_case_ =   True
            else:
                  for key, mapped_key in MAPPING.items():
                        # mapped_key = "speecht5." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key

                        if "*" in key:
                              snake_case_ ,					snake_case_ =   key.split('''.*.'''      )
                              if prefix in name and suffix in name:
                                    snake_case_ =   suffix

                # if key in name or key.split("w2v_model.")[-1] == name.split(".")[0]:
                        if key in name:
                              snake_case_ =   True
                              if "*" in mapped_key:
                                    snake_case_ =   name.split(lowerCAmelCase__      )[0].split('''.'''      )[-2]
                                    snake_case_ =   mapped_key.replace('''*'''      ,      lowerCAmelCase__      )
                              if "weight_g" in name:
                                    snake_case_ =   '''weight_g'''
                              elif "weight_v" in name:
                                    snake_case_ =   '''weight_v'''
                              elif "bias" in name:
                                    snake_case_ =   '''bias'''
                              elif "weight" in name:
                                    snake_case_ =   '''weight'''
                              elif "running_mean" in name:
                                    snake_case_ =   '''running_mean'''
                              elif "running_var" in name:
                                    snake_case_ =   '''running_var'''
                              elif "num_batches_tracked" in name:
                                    snake_case_ =   '''num_batches_tracked'''
                              else:
                                    snake_case_ =   None
                              set_recursively(lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      )
                        continue
            if not is_used:
                  unused_weights.append(lowerCAmelCase__      )

      logger.warning(f'''Unused weights: {unused_weights}'''      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						str:




      """simple docstring"""

      snake_case_ =   full_name.split('''conv_layers.'''      )[-1]
      snake_case_ =   name.split('''.'''      )
      snake_case_ =   int(items[0]      )
      snake_case_ =   int(items[1]      )

      if type_id == 0:
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                        raise ValueError(
                            f'''{full_name} has size {value.shape}, but'''
                            f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''      )
                  snake_case_ =   value
                  logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''      )
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                        raise ValueError(
                            f'''{full_name} has size {value.shape}, but'''
                            f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''      )
                  snake_case_ =   value
                  logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''      )
      elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                        raise ValueError(
                            f'''{full_name} has size {value.shape}, but'''
                            f''' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'''      )
                  snake_case_ =   value
                  logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''      )
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                        raise ValueError(
                            f'''{full_name} has size {value.shape}, but'''
                            f''' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'''      )
                  snake_case_ =   value
                  logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''      )
      else:
            unused_weights.append(lowerCAmelCase__      )






@torch.no_grad()
def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=None      ,      SCREAMING_SNAKE_CASE=None      ,      SCREAMING_SNAKE_CASE=None      ,      )->						int:




      """simple docstring"""

      if config_path is not None:
            snake_case_ =   SpeechTaConfig.from_pretrained(lowerCAmelCase__      )
      else:
            snake_case_ =   SpeechTaConfig()

      if task == "s2t":
            snake_case_ =   config.max_text_positions
            snake_case_ =   SpeechTaForSpeechToText(lowerCAmelCase__      )
      elif task == "t2s":
            snake_case_ =   1876
            snake_case_ =   600
            snake_case_ =   config.max_speech_positions
            snake_case_ =   SpeechTaForTextToSpeech(lowerCAmelCase__      )
      elif task == "s2s":
            snake_case_ =   1876
            snake_case_ =   config.max_speech_positions
            snake_case_ =   SpeechTaForSpeechToSpeech(lowerCAmelCase__      )
      else:
            raise ValueError(f'''Unknown task name: {task}'''      )

      if vocab_path:
            snake_case_ =   SpeechTaTokenizer(lowerCAmelCase__      ,      model_max_length=config.max_text_positions      )

            # Mask token behaves like a normal word, i.e. include the space before it
            snake_case_ =   AddedToken('''<mask>'''      ,      lstrip=lowerCAmelCase__      ,      rstrip=lowerCAmelCase__      )
            snake_case_ =   mask_token
            tokenizer.add_special_tokens({'''mask_token''': mask_token}      )
            tokenizer.add_tokens(['''<ctc_blank>''']      )

      snake_case_ =   SpeechTaFeatureExtractor()
      snake_case_ =   SpeechTaProcessor(tokenizer=lowerCAmelCase__      ,      feature_extractor=lowerCAmelCase__      )
      processor.save_pretrained(lowerCAmelCase__      )

      snake_case_ =   torch.load(lowerCAmelCase__      )
      recursively_load_weights(fairseq_checkpoint['''model''']      ,      lowerCAmelCase__      ,      lowerCAmelCase__      )

      model.save_pretrained(lowerCAmelCase__      )

      if repo_id:
            print('''Pushing to the hub...'''      )
            processor.push_to_hub(lowerCAmelCase__      )
            model.push_to_hub(lowerCAmelCase__      )


if __name__ == "__main__":
   UpperCAmelCase    =					argparse.ArgumentParser()
   parser.add_argument(
       """--task""",
       default="""s2t""",
       type=str,
       help="""Type of the SpeechT5 model you\'d like to convert. Should be one of \'s2t\', \'t2s\', \'s2s\'.""",
   )
   parser.add_argument("""--checkpoint_path""", required=True, default=None, type=str, help="""Path to fairseq checkpoint""")
   parser.add_argument("""--vocab_path""", default=None, type=str, help="""Path to SentencePiece model""")
   parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
   parser.add_argument(
       """--pytorch_dump_folder_path""", required=True, default=None, type=str, help="""Path to the output PyTorch model."""
   )
   parser.add_argument(
       """--push_to_hub""", default=None, type=str, help="""Where to upload the converted model on the 🤗 hub."""
   )

   UpperCAmelCase    =					parser.parse_args()
   convert_speechta_checkpoint(
       args.task,
       args.checkpoint_path,
       args.pytorch_dump_folder_path,
       args.config_path,
       args.vocab_path,
       args.push_to_hub,
   )



import unittest

import numpy as np

from diffusers import LMSDiscreteScheduler, OnnxStableDiffusionInpaintPipeline
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
   import onnxruntime as ort







class lowerCAmelCase_		(				lowerCamelCase__      ,		unittest.TestCase ):

 '''simple docstring'''

 pass







@nightly
@require_onnxruntime
@require_torch_gpu
class lowerCAmelCase_		(				unittest.TestCase ):

 '''simple docstring'''

 @property
 def 		UpperCamelCase__  (			self      ):
       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )



 @property
 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   ort.SessionOptions()
       snake_case_ =   False
       return options



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/in_paint/overture-creations-5sI6fQgYIuo.png'''      )
       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/in_paint/overture-creations-5sI6fQgYIuo_mask.png'''      )
       snake_case_ =   OnnxStableDiffusionInpaintPipeline.from_pretrained(
           '''runwayml/stable-diffusion-inpainting'''	,   revision='''onnx'''	,   safety_checker=_UpperCAmelCase	,   feature_extractor=_UpperCAmelCase	,   provider=self.gpu_provider	,   sess_options=self.gpu_options	,   )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase      )

       snake_case_ =   '''A red cat sitting on a park bench'''

       snake_case_ =   np.random.RandomState(0      )
       snake_case_ =   pipe(
           prompt=_UpperCAmelCase	,   image=_UpperCAmelCase	,   mask_image=_UpperCAmelCase	,   guidance_scale=7.5	,   num_inference_steps=10	,   generator=_UpperCAmelCase	,   output_type='''np'''	,   )
       snake_case_ =   output.images
       snake_case_ =   images[0, 2_55:2_58, 2_55:2_58, -1]

       assert images.shape == (1, 5_12, 5_12, 3)
       snake_case_ =   np.array([0.2_514, 0.3_007, 0.3_517, 0.1_790, 0.2_382, 0.3_167, 0.1_944, 0.2_273, 0.2_464]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-3





 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/in_paint/overture-creations-5sI6fQgYIuo.png'''      )
       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/in_paint/overture-creations-5sI6fQgYIuo_mask.png'''      )
       snake_case_ =   LMSDiscreteScheduler.from_pretrained(
           '''runwayml/stable-diffusion-inpainting'''	,   subfolder='''scheduler'''	,   revision='''onnx'''      )
       snake_case_ =   OnnxStableDiffusionInpaintPipeline.from_pretrained(
           '''runwayml/stable-diffusion-inpainting'''	,   revision='''onnx'''	,   scheduler=_UpperCAmelCase	,   safety_checker=_UpperCAmelCase	,   feature_extractor=_UpperCAmelCase	,   provider=self.gpu_provider	,   sess_options=self.gpu_options	,   )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase      )

       snake_case_ =   '''A red cat sitting on a park bench'''

       snake_case_ =   np.random.RandomState(0      )
       snake_case_ =   pipe(
           prompt=_UpperCAmelCase	,   image=_UpperCAmelCase	,   mask_image=_UpperCAmelCase	,   guidance_scale=7.5	,   num_inference_steps=20	,   generator=_UpperCAmelCase	,   output_type='''np'''	,   )
       snake_case_ =   output.images
       snake_case_ =   images[0, 2_55:2_58, 2_55:2_58, -1]

       assert images.shape == (1, 5_12, 5_12, 3)
       snake_case_ =   np.array([0.0_086, 0.0_077, 0.0_083, 0.0_093, 0.0_107, 0.0_139, 0.0_094, 0.0_097, 0.0_125]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-3