





'''simple docstring'''





import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu



class    _a     (				nn.Module							):






							"""simple docstring"""

							def __init__(      self				:     Optional[int]    )	->Any:
													super().__init__()
													SCREAMING_SNAKE_CASE__     :				Any   =					nn.Linear(3							,	4    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					nn.BatchNormad(4    )
													SCREAMING_SNAKE_CASE__     :				Any   =					nn.Linear(4							,	5    )






							def      A_  (      self				:     List[Any]							,	a				:     Dict    )	->Any:
													return self.lineara(self.batchnorm(self.lineara(a    )    )    )







class    _a     (				lowercase__							):






							"""simple docstring"""

							def      A_  (      self				:     Optional[Any]							,	a				:     Union[str, Any]							,	*a				:     Union[str, Any]							,	**a				:     List[Any]    )	->Any:
													return (args[0] + 1,) + args[1:], kwargs







class    _a     (				lowercase__							):






							"""simple docstring"""

							def      A_  (      self				:     str							,	a				:     Tuple							,	a				:     Tuple    )	->str:
													return output + 1







class    _a     (				unittest.TestCase							):






							"""simple docstring"""

							def      A_  (      self				:     Any    )	->str:
													SCREAMING_SNAKE_CASE__     :				List[str]   =					ModelForTest()
													SCREAMING_SNAKE_CASE__     :				Dict   =					ModelHook()

													add_hook_to_module(a							,	a    )
													self.assertEqual(test_model._hf_hook							,	a    )
													self.assertTrue(hasattr(a							,	"_old_forward"    )    )

													# Check adding the hook did not change the name or the signature
													self.assertEqual(test_model.forward.__name__							,	"forward"    )
													self.assertListEqual(list(inspect.signature(test_model.forward    ).parameters    )							,	["x"]    )

													remove_hook_from_module(a    )
													self.assertFalse(hasattr(a							,	"_hf_hook"    )    )
													self.assertFalse(hasattr(a							,	"_old_forward"    )    )

							def      A_  (      self				:     str    )	->List[Any]:
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					ModelForTest()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					ModelHook()

													add_hook_to_module(a							,	a    )
													add_hook_to_module(a							,	a							,	append=a    )

													self.assertEqual(isinstance(test_model._hf_hook							,	a    )							,	a    )
													self.assertEqual(len(test_model._hf_hook.hooks    )							,	2    )
													self.assertTrue(hasattr(a							,	"_old_forward"    )    )

													# Check adding the hook did not change the name or the signature
													self.assertEqual(test_model.forward.__name__							,	"forward"    )
													self.assertListEqual(list(inspect.signature(test_model.forward    ).parameters    )							,	["x"]    )

													remove_hook_from_module(a    )
													self.assertFalse(hasattr(a							,	"_hf_hook"    )    )
													self.assertFalse(hasattr(a							,	"_old_forward"    )    )

							def      A_  (      self				:     List[str]    )	->List[Any]:
													SCREAMING_SNAKE_CASE__     :				Any   =					ModelForTest()
													SCREAMING_SNAKE_CASE__     :				Any   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Any   =					test_model(x + 1    )
													SCREAMING_SNAKE_CASE__     :				int   =					test_model(x + 2    )

													SCREAMING_SNAKE_CASE__     :				List[str]   =					PreForwardHook()
													add_hook_to_module(a							,	a    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					test_model(a    )
													self.assertTrue(torch.allclose(a							,	a							,	atol=1E-5    )    )

													# Attaching a hook to a model when it already has one replaces, does not chain
													SCREAMING_SNAKE_CASE__     :				Dict   =					PreForwardHook()
													add_hook_to_module(a							,	a    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					test_model(a    )
													self.assertTrue(torch.allclose(a							,	a							,	atol=1E-5    )    )

													# You need to use the sequential hook to chain two or more hooks
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					SequentialHook(PreForwardHook()							,	PreForwardHook()    )
													add_hook_to_module(a							,	a    )

													SCREAMING_SNAKE_CASE__     :				str   =					test_model(a    )
													assert torch.allclose(a							,	a							,	atol=1E-5    )

							def      A_  (      self				:     Optional[Any]    )	->Tuple:
													SCREAMING_SNAKE_CASE__     :				int   =					ModelForTest()
													SCREAMING_SNAKE_CASE__     :				Tuple   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					test_model(a    )

													SCREAMING_SNAKE_CASE__     :				str   =					PostForwardHook()
													add_hook_to_module(a							,	a    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					test_model(a    )
													self.assertTrue(torch.allclose(a							,	output + 1							,	atol=1E-5    )    )

													# Attaching a hook to a model when it already has one replaces, does not chain
													SCREAMING_SNAKE_CASE__     :				Tuple   =					PostForwardHook()
													add_hook_to_module(a							,	a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					test_model(a    )
													self.assertTrue(torch.allclose(a							,	output + 1							,	atol=1E-5    )    )

													# You need to use the sequential hook to chain two or more hooks
													SCREAMING_SNAKE_CASE__     :				List[str]   =					SequentialHook(PostForwardHook()							,	PostForwardHook()    )
													add_hook_to_module(a							,	a    )

													SCREAMING_SNAKE_CASE__     :				Tuple   =					test_model(a    )
													assert torch.allclose(a							,	output + 2							,	atol=1E-5    )

							def      A_  (      self				:     int    )	->Optional[int]:
													SCREAMING_SNAKE_CASE__     :				str   =					ModelForTest()
													SCREAMING_SNAKE_CASE__     :				Tuple   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				int   =					test_model(a    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					PostForwardHook()
													add_hook_to_module(a							,	a    )
													SCREAMING_SNAKE_CASE__     :				int   =					test_model(a    )
													self.assertTrue(torch.allclose(a							,	output + 1    )    )
													self.assertTrue(outputa.requires_grad    )

													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					True
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					test_model(a    )
													self.assertFalse(outputa.requires_grad    )

							@require_multi_gpu
							def      A_  (      self				:     int    )	->Dict:
													SCREAMING_SNAKE_CASE__     :				Any   =					ModelForTest()
													# Everything is on CPU
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# This will move each submodule on different devices
													add_hook_to_module(model.lineara							,	AlignDevicesHook(execution_device=0    )    )
													add_hook_to_module(model.batchnorm							,	AlignDevicesHook(execution_device=0    )    )
													add_hook_to_module(model.lineara							,	AlignDevicesHook(execution_device=1    )    )

													self.assertEqual(model.lineara.weight.device							,	torch.device(0    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device(0    )    )
													self.assertEqual(model.batchnorm.running_mean.device							,	torch.device(0    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device(1    )    )

													# We can still make a forward pass. The input does not need to be on any particular device
													SCREAMING_SNAKE_CASE__     :				Any   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					model(a    )
													self.assertEqual(output.device							,	torch.device(1    )    )

													# We can add a general hook to put back output on same device as input.
													add_hook_to_module(a							,	AlignDevicesHook(io_same_device=a    )    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					torch.randn(2							,	3    ).to(0    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					model(a    )
													self.assertEqual(output.device							,	torch.device(0    )    )

							def      A_  (      self				:     Tuple    )	->List[Any]:
													SCREAMING_SNAKE_CASE__     :				Tuple   =					ModelForTest()

													# Everything is on CPU
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# This will move each submodule on different devices
													SCREAMING_SNAKE_CASE__     :				int   =					{"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

													add_hook_to_module(model.lineara							,	AlignDevicesHook(**a    )    )
													add_hook_to_module(model.batchnorm							,	AlignDevicesHook(**a    )    )
													add_hook_to_module(model.lineara							,	AlignDevicesHook(**a    )    )

													# Parameters have been offloaded, so on the meta device
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													# Buffers are not included in the offload by default, so are on the execution device
													SCREAMING_SNAKE_CASE__     :				str   =					torch.device(hook_kwargs["execution_device"]    )
													self.assertEqual(model.batchnorm.running_mean.device							,	a    )

													SCREAMING_SNAKE_CASE__     :				int   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				str   =					model(a    )
													self.assertEqual(output.device							,	a    )

													# Removing hooks loads back the weights in the model.
													remove_hook_from_module(model.lineara    )
													remove_hook_from_module(model.batchnorm    )
													remove_hook_from_module(model.lineara    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# Now test with buffers included in the offload
													SCREAMING_SNAKE_CASE__     :				Tuple   =					{
													    "execution_device": 0 if torch.cuda.is_available() else "cpu",
													    "offload": True,
													    "offload_buffers": True,
													}

													add_hook_to_module(model.lineara							,	AlignDevicesHook(**a    )    )
													add_hook_to_module(model.batchnorm							,	AlignDevicesHook(**a    )    )
													add_hook_to_module(model.lineara							,	AlignDevicesHook(**a    )    )

													# Parameters have been offloaded, so on the meta device, buffers included
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.running_mean.device							,	torch.device("meta"    )    )

													SCREAMING_SNAKE_CASE__     :				int   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					model(a    )
													self.assertEqual(output.device							,	a    )

													# Removing hooks loads back the weights in the model.
													remove_hook_from_module(model.lineara    )
													remove_hook_from_module(model.batchnorm    )
													remove_hook_from_module(model.lineara    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

							def      A_  (      self				:     Dict    )	->Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :				int   =					ModelForTest()

													# Everything is on CPU
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# This will move each submodule on different devices
													SCREAMING_SNAKE_CASE__     :				Any   =					0 if torch.cuda.is_available() else "cpu"
													attach_align_device_hook(a							,	execution_device=a							,	offload=a    )

													# Parameters have been offloaded, so on the meta device
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													# Buffers are not included in the offload by default, so are on the execution device
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					torch.device(a    )
													self.assertEqual(model.batchnorm.running_mean.device							,	a    )

													SCREAMING_SNAKE_CASE__     :				Any   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					model(a    )
													self.assertEqual(output.device							,	a    )

													# Removing hooks loads back the weights in the model.
													remove_hook_from_submodules(a    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# Now test with buffers included in the offload
													attach_align_device_hook(a							,	execution_device=a							,	offload=a							,	offload_buffers=a    )

													# Parameters have been offloaded, so on the meta device, buffers included
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.running_mean.device							,	torch.device("meta"    )    )

													SCREAMING_SNAKE_CASE__     :				Tuple   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Any   =					model(a    )
													self.assertEqual(output.device							,	a    )

													# Removing hooks loads back the weights in the model.
													remove_hook_from_submodules(a    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )






							def      A_  (      self				:     Tuple    )	->List[str]:
													SCREAMING_SNAKE_CASE__     :				Tuple   =					ModelForTest()

													# Everything is on CPU
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# This will move each submodule on different devices
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					0 if torch.cuda.is_available() else "cpu"
													attach_align_device_hook(
													    a							,	execution_device=a							,	offload=a							,	weights_map=model.state_dict()    )

													# Parameters have been offloaded, so on the meta device
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													# Buffers are not included in the offload by default, so are on the execution device
													SCREAMING_SNAKE_CASE__     :				str   =					torch.device(a    )
													self.assertEqual(model.batchnorm.running_mean.device							,	a    )

													SCREAMING_SNAKE_CASE__     :				Any   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				str   =					model(a    )
													self.assertEqual(output.device							,	a    )

													# Removing hooks loads back the weights in the model.
													remove_hook_from_submodules(a    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )

													# Now test with buffers included in the offload
													attach_align_device_hook(
													    a							,	execution_device=a							,	offload=a							,	weights_map=model.state_dict()							,	offload_buffers=a							,	)

													# Parameters have been offloaded, so on the meta device, buffers included
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("meta"    )    )
													self.assertEqual(model.batchnorm.running_mean.device							,	torch.device("meta"    )    )

													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					torch.randn(2							,	3    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					model(a    )
													self.assertEqual(output.device							,	a    )

													# Removing hooks loads back the weights in the model.
													remove_hook_from_submodules(a    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.batchnorm.weight.device							,	torch.device("cpu"    )    )
													self.assertEqual(model.lineara.weight.device							,	torch.device("cpu"    )    )






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






							"""simple docstring"""

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

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

													SCREAMING_SNAKE_CASE__     :				Dict   =					"bird"
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.device_count()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe.prepare_text_inputs([prompts] * num_samples    )

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

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.random.PRNGKey(0    )
													SCREAMING_SNAKE_CASE__     :				int   =					jax.random.split(a							,	jax.device_count()    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					replicate(a    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					shard(a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					shard(a    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					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, 7_68, 5_12, 3)

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					images[0, 2_53:2_56, 2_53:2_56, -1]

													SCREAMING_SNAKE_CASE__     :				Tuple   =					jnp.asarray(jax.device_get(image_slice.flatten()    )    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					jnp.array(
													    [0.16_7969, 0.11_6699, 0.08_1543, 0.15_4297, 0.13_2812, 0.10_8887, 0.16_9922, 0.16_9922, 0.20_5078]    )
													print(f"""output_slice: {output_slice}"""    )
													assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2






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

													SCREAMING_SNAKE_CASE__     :				Any   =					"Chef in the kitchen"
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jax.device_count()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe.prepare_text_inputs([prompts] * num_samples    )

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

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

													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					replicate(a    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					shard(a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					shard(a    )

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					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, 7_68, 5_12, 3)

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
													SCREAMING_SNAKE_CASE__     :				str   =					images[0, 2_53:2_56, 2_53:2_56, -1]

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.asarray(jax.device_get(image_slice.flatten()    )    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.array(
													    [[0.27_1484, 0.26_1719, 0.27_5391, 0.27_7344, 0.27_9297, 0.29_1016, 0.29_4922, 0.30_2734, 0.30_2734]]    )
													print(f"""output_slice: {output_slice}"""    )
													assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2