





import unittest

from transformers import DebertaVaTokenizer, DebertaVaTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


snake_case_           =      get_tests_dir('fixtures/spiece.model')

@require_sentencepiece
@require_tokenizers
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    Any							=   DebertaVaTokenizer
      A_						:    List[Any]							=   DebertaVaTokenizerFast
      A_						:    int							=   True
      A_						:    List[Any]							=   True


      def        a			(self	:				str				):





         """simple docstring"""



         super().setUp()

         # We have a SentencePiece fixture for testing
         __snake_case								=       DebertaVaTokenizer(a__    ,      unk_token='''<unk>'''				)
         tokenizer.save_pretrained(self.tmpdirname				)


      def        a			(self	:				List[Any]    ,      a__	:				List[str]				):





         """simple docstring"""



         __snake_case								=       '''this is a test'''
         __snake_case								=       '''this is a test'''
         return input_text, output_text


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





         """simple docstring"""



         __snake_case								=       '''<pad>'''
         __snake_case								=       0

         self.assertEqual(self.get_tokenizer()._convert_token_to_id(a__				)    ,      a__				)
         self.assertEqual(self.get_tokenizer()._convert_id_to_token(a__				)    ,      a__				)


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





         """simple docstring"""



         __snake_case								=       list(self.get_tokenizer().get_vocab().keys()				)
         self.assertEqual(vocab_keys[0]    ,      '''<pad>'''				)
         self.assertEqual(vocab_keys[1]    ,      '''<unk>'''				)
         self.assertEqual(vocab_keys[-1]    ,      '''[PAD]'''				)
         self.assertEqual(len(a__				)    ,      3_0001				)


      def        a			(self	:				Dict				):





         """simple docstring"""



         self.assertEqual(self.get_tokenizer().vocab_size    ,      3_0000				)


      def        a			(self	:				str				):





         """simple docstring"""



         __snake_case								=       ''' \tHeLLo!how  \n Are yoU?  '''
         __snake_case								=       ['''▁hello''', '''!''', '''how''', '''▁are''', '''▁you''', '''?''']
         # fmt: on

         __snake_case								=       DebertaVaTokenizer(a__    ,      do_lower_case=a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       DebertaVaTokenizerFast(a__    ,      do_lower_case=a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)


      @unittest.skip('''There is an inconsistency between slow and fast tokenizer due to a bug in the fast one.'''				)
      def        a			(self	:				Any				):





         """simple docstring"""



         pass


      @unittest.skip('''There is an inconsistency between slow and fast tokenizer due to a bug in the fast one.'''				)
      def        a			(self	:				Dict				):





         """simple docstring"""



         pass


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       '''I was born in 92000, and this is falsé.'''
         __snake_case								=       ['''▁''', '''<unk>''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', '''▁''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''<unk>''', '''▁''', '''.''', ]
         # fmt: on

         __snake_case								=       DebertaVaTokenizer(a__    ,      split_by_punct=a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       DebertaVaTokenizerFast(a__    ,      split_by_punct=a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       '''I was born in 92000, and this is falsé.'''
         __snake_case								=       ['''▁i''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', '''▁''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''<unk>''', '''▁''', '''.''', ]
         # fmt: on

         __snake_case								=       DebertaVaTokenizer(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)
         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       DebertaVaTokenizerFast(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)
         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       '''I was born in 92000, and this is falsé.'''
         __snake_case								=       ['''▁i''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''<unk>''', '''.''', ]
         # fmt: on

         __snake_case								=       DebertaVaTokenizer(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       DebertaVaTokenizerFast(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case								=       '''I was born in 92000, and this is falsé.'''
         __snake_case								=       ['''▁''', '''<unk>''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', '''▁''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''<unk>''', '''▁''', '''.''', ]
         # fmt: on

         __snake_case								=       DebertaVaTokenizer(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       DebertaVaTokenizerFast(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       ''' \tHeLLo!how  \n Are yoU?  '''
         __snake_case								=       ['''▁''', '''<unk>''', '''e''', '''<unk>''', '''o''', '''!''', '''how''', '''▁''', '''<unk>''', '''re''', '''▁yo''', '''<unk>''', '''?''']
         # fmt: on

         __snake_case								=       DebertaVaTokenizer(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       DebertaVaTokenizerFast(a__    ,      do_lower_case=a__    ,      split_by_punct=a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)

         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				Any				):





         """simple docstring"""



         __snake_case								=       self.get_tokenizer()
         __snake_case								=       self.get_rust_tokenizer()

         __snake_case								=       '''I was born in 92000, and this is falsé.'''

         __snake_case								=       tokenizer.convert_ids_to_tokens(tokenizer.encode(a__    ,      add_special_tokens=a__				)				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)				)
         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       tokenizer.encode(a__    ,      add_special_tokens=a__				)
         __snake_case								=       rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)
         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       self.get_rust_tokenizer()
         __snake_case								=       tokenizer.encode(a__				)
         __snake_case								=       rust_tokenizer.encode(a__				)
         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				str				):





         """simple docstring"""



         __snake_case								=       '''This is a test'''
         __snake_case								=       [13, 1, 4398, 25, 21, 1289]
         __snake_case								=       ['''▁''', '''T''', '''his''', '''▁is''', '''▁a''', '''▁test''']
         __snake_case								=       ['''▁''', '''<unk>''', '''his''', '''▁is''', '''▁a''', '''▁test''']

         __snake_case								=       DebertaVaTokenizer(a__    ,      keep_accents=a__				)
         __snake_case								=       DebertaVaTokenizerFast(a__    ,      keep_accents=a__				)

         __snake_case								=       tokenizer.encode(a__    ,      add_special_tokens=a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       tokenizer.tokenize(a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(a__				)
         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       rust_tokenizer.tokenize(a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(a__				)
         self.assertListEqual(a__    ,      a__				)

         # fmt: off
         __snake_case								=       '''I was born in 92000, and this is falsé.'''
         __snake_case								=       [13, 1, 23, 386, 19, 561, 3050, 15, 17, 48, 25, 8256, 18, 1, 9]
         __snake_case								=       ['''▁''', '''I''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''é''', '''.''', ]
         __snake_case								=       ['''▁''', '''<unk>''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''<unk>''', '''.''', ]
         # fmt: on

         __snake_case								=       tokenizer.encode(a__    ,      add_special_tokens=a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       tokenizer.tokenize(a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       tokenizer.convert_ids_to_tokens(a__				)
         self.assertListEqual(a__    ,      a__				)

         __snake_case								=       rust_tokenizer.encode(a__    ,      add_special_tokens=a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       rust_tokenizer.tokenize(a__				)
         self.assertListEqual(a__    ,      a__				)
         __snake_case								=       rust_tokenizer.convert_ids_to_tokens(a__				)
         self.assertListEqual(a__    ,      a__				)


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         __snake_case								=       DebertaVaTokenizer(a__				)

         __snake_case								=       tokenizer.encode('''sequence builders'''				)
         __snake_case								=       tokenizer.encode('''multi-sequence build'''				)

         __snake_case								=       tokenizer.build_inputs_with_special_tokens(a__				)
         __snake_case								=       tokenizer.build_inputs_with_special_tokens(a__    ,      a__				)

         self.assertEqual([tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]    ,      a__				)
         self.assertEqual(
             [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + text_a + [tokenizer.sep_token_id]    ,      a__    ,      )


      @slow
      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       {'''input_ids''': [[1, 3_9867, 36, 1_9390, 486, 27, 3_5052, 8_1436, 18, 6_0685, 1225, 7, 3_5052, 8_1436, 18, 9367, 1_6899, 18, 1_5937, 53, 594, 773, 18, 1_6287, 3_0465, 36, 1_5937, 6, 4_1139, 38, 3_6979, 6_0763, 191, 6, 3_4132, 99, 6, 5_0538, 390, 4_3230, 6, 3_4132, 2779, 2_0850, 14, 699, 1072, 1194, 36, 382, 1_0901, 53, 7, 699, 1072, 2084, 36, 2_0422, 630, 53, 19, 105, 3049, 1896, 1053, 1_6899, 1506, 11, 3_7978, 4243, 7, 1237, 3_1869, 200, 1_6566, 654, 6, 3_5052, 8_1436, 7, 5_5630, 1_3593, 4, 2], [1, 26, 1_5011, 13, 667, 8, 1053, 18, 2_3611, 1237, 7_2356, 1_2820, 34, 10_4134, 1209, 35, 1_3313, 6627, 21, 202, 347, 7, 164, 2399, 11, 46, 4485, 4, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 5, 1232, 2864, 1_5785, 1_4951, 105, 5, 8581, 1250, 4, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''token_type_ids''': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''attention_mask''': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
         # fmt: on

         self.tokenizer_integration_test_util(
             expected_encoding=a__    ,      model_name='''microsoft/deberta-v2-xlarge'''    ,      revision='''ad6e42c1532ddf3a15c39246b63f5559d558b670'''    ,      )








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  ):


      def        a			(self	:				int				):





         """simple docstring"""



         super().tearDown()
         gc.collect()


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       FlaxControlNetModel.from_pretrained(
             '''lllyasviel/sd-controlnet-canny'''    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case	,     __snake_case								=       FlaxStableDiffusionControlNetPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''    ,      controlnet=a__    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case								=       controlnet_params

         __snake_case								=       '''bird'''
         __snake_case								=       jax.device_count()
         __snake_case								=       pipe.prepare_text_inputs([prompts] * num_samples				)

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

         __snake_case								=       jax.random.PRNGKey(0				)
         __snake_case								=       jax.random.split(a__    ,      jax.device_count()				)

         __snake_case								=       replicate(a__				)
         __snake_case								=       shard(a__				)
         __snake_case								=       shard(a__				)

         __snake_case								=       pipe(
             prompt_ids=a__    ,      image=a__    ,      params=a__    ,      prng_seed=a__    ,      num_inference_steps=50    ,      jit=a__    ,      ).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

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

         __snake_case								=       jnp.asarray(jax.device_get(image_slice.flatten()				)				)
         __snake_case								=       jnp.array(
             [0.1_6_7_9_6_9, 0.1_1_6_6_9_9, 0.0_8_1_5_4_3, 0.1_5_4_2_9_7, 0.1_3_2_8_1_2, 0.1_0_8_8_8_7, 0.1_6_9_9_2_2, 0.1_6_9_9_2_2, 0.2_0_5_0_7_8]				)
         print(f"""output_slice: {output_slice}"""				)
         assert jnp.abs(output_slice - expected_slice				).max() < 1E-2


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       FlaxControlNetModel.from_pretrained(
             '''lllyasviel/sd-controlnet-openpose'''    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case	,     __snake_case								=       FlaxStableDiffusionControlNetPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''    ,      controlnet=a__    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case								=       controlnet_params

         __snake_case								=       '''Chef in the kitchen'''
         __snake_case								=       jax.device_count()
         __snake_case								=       pipe.prepare_text_inputs([prompts] * num_samples				)

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

         __snake_case								=       jax.random.PRNGKey(0				)
         __snake_case								=       jax.random.split(a__    ,      jax.device_count()				)

         __snake_case								=       replicate(a__				)
         __snake_case								=       shard(a__				)
         __snake_case								=       shard(a__				)

         __snake_case								=       pipe(
             prompt_ids=a__    ,      image=a__    ,      params=a__    ,      prng_seed=a__    ,      num_inference_steps=50    ,      jit=a__    ,      ).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

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

         __snake_case								=       jnp.asarray(jax.device_get(image_slice.flatten()				)				)
         __snake_case								=       jnp.array(
             [[0.2_7_1_4_8_4, 0.2_6_1_7_1_9, 0.2_7_5_3_9_1, 0.2_7_7_3_4_4, 0.2_7_9_2_9_7, 0.2_9_1_0_1_6, 0.2_9_4_9_2_2, 0.3_0_2_7_3_4, 0.3_0_2_7_3_4]]				)
         print(f"""output_slice: {output_slice}"""				)
         assert jnp.abs(output_slice - expected_slice				).max() < 1E-2

