


'''simple docstring'''



import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
		import onnxruntime as ort


class 		_snake_case					(     a__       ,   unittest.TestCase	):
					lowerCAmelCase      :List[str]   =					'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''




					def 				snake_case__						(      self							,							_lowerCamelCase=0):
								UpperCAmelCase__  :					List[Any]        =     np.random.RandomState(_lowerCamelCase)
								UpperCAmelCase__  :					str        =     {
								    """prompt""": """A painting of a squirrel eating a burger""",
								    """generator""": generator,
								    """num_inference_steps""": 2,
								    """guidance_scale""": 7.5,
								    """output_type""": """numpy""",
								}
								return inputs




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					str        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					List[Any]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Any        =     pipe(**_lowerCamelCase).images
								UpperCAmelCase__  :					List[Any]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 128, 128, 3)
								UpperCAmelCase__  :					List[Any]        =     np.array([0.65072, 0.58492, 0.48219, 0.55521, 0.53180, 0.55939, 0.50697, 0.39800, 0.46455])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[Any]        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								UpperCAmelCase__  :					Any        =     PNDMScheduler.from_config(pipe.scheduler.config							,							skip_prk_steps=_lowerCamelCase)
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Dict        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Optional[int]        =     pipe(**_lowerCamelCase).images
								UpperCAmelCase__  :					Optional[int]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 128, 128, 3)
								UpperCAmelCase__  :					Optional[int]        =     np.array([0.65863, 0.59425, 0.49326, 0.56313, 0.53875, 0.56627, 0.51065, 0.39777, 0.46330])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Any        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								UpperCAmelCase__  :					Tuple        =     LMSDiscreteScheduler.from_config(pipe.scheduler.config)
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Optional[int]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					str        =     pipe(**_lowerCamelCase).images
								UpperCAmelCase__  :					Optional[Any]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 128, 128, 3)
								UpperCAmelCase__  :					Optional[int]        =     np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Dict        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								UpperCAmelCase__  :					List[Any]        =     EulerDiscreteScheduler.from_config(pipe.scheduler.config)
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Union[str, Any]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Dict        =     pipe(**_lowerCamelCase).images
								UpperCAmelCase__  :					Optional[int]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 128, 128, 3)
								UpperCAmelCase__  :					Dict        =     np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					int        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								UpperCAmelCase__  :					Optional[Any]        =     EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Any        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Union[str, Any]        =     pipe(**_lowerCamelCase).images
								UpperCAmelCase__  :					Optional[Any]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 128, 128, 3)
								UpperCAmelCase__  :					Tuple        =     np.array([0.53817, 0.60812, 0.47384, 0.49530, 0.51894, 0.49814, 0.47984, 0.38958, 0.44271])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[Any]        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								UpperCAmelCase__  :					int        =     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					List[str]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					List[Any]        =     pipe(**_lowerCamelCase).images
								UpperCAmelCase__  :					Optional[Any]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 128, 128, 3)
								UpperCAmelCase__  :					Tuple        =     np.array([0.53895, 0.60808, 0.47933, 0.49608, 0.51886, 0.49950, 0.48053, 0.38957, 0.44200])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[int]        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					int        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Optional[Any]        =     3 * [inputs["""prompt"""]]

								# forward
								UpperCAmelCase__  :					Union[str, Any]        =     pipe(**_lowerCamelCase)
								UpperCAmelCase__  :					Union[str, Any]        =     output.images[0, -3:, -3:, -1]

								UpperCAmelCase__  :					Optional[Any]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					int        =     3 * [inputs.pop("""prompt""")]

								UpperCAmelCase__  :					Optional[Any]        =     pipe.tokenizer(
								    _lowerCamelCase							,							padding="""max_length"""							,							max_length=pipe.tokenizer.model_max_length							,							truncation=_lowerCamelCase							,							return_tensors="""np"""							,							)
								UpperCAmelCase__  :					List[Any]        =     text_inputs["""input_ids"""]

								UpperCAmelCase__  :					str        =     pipe.text_encoder(input_ids=text_inputs.astype(np.intaa))[0]

								UpperCAmelCase__  :					int        =     prompt_embeds

								# forward
								UpperCAmelCase__  :					List[str]        =     pipe(**_lowerCamelCase)
								UpperCAmelCase__  :					str        =     output.images[0, -3:, -3:, -1]

								assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()).max() < 1e-4




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					str        =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,							provider="""CPUExecutionProvider""")
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Union[str, Any]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Any        =     3 * ["""this is a negative prompt"""]
								UpperCAmelCase__  :					Any        =     negative_prompt
								UpperCAmelCase__  :					Optional[Any]        =     3 * [inputs["""prompt"""]]

								# forward
								UpperCAmelCase__  :					Optional[int]        =     pipe(**_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     output.images[0, -3:, -3:, -1]

								UpperCAmelCase__  :					Optional[Any]        =     self.get_dummy_inputs()
								UpperCAmelCase__  :					Union[str, Any]        =     3 * [inputs.pop("""prompt""")]

								UpperCAmelCase__  :					List[str]        =     []
								for p in [prompt, negative_prompt]:
											UpperCAmelCase__  :					Tuple        =     pipe.tokenizer(
											    _lowerCamelCase							,							padding="""max_length"""							,							max_length=pipe.tokenizer.model_max_length							,							truncation=_lowerCamelCase							,							return_tensors="""np"""							,							)
											UpperCAmelCase__  :					Optional[int]        =     text_inputs["""input_ids"""]

											embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa))[0])

								UpperCAmelCase__  :					Optional[Any]        =     embeds

								# forward
								UpperCAmelCase__  :					Any        =     pipe(**_lowerCamelCase)
								UpperCAmelCase__  :					Any        =     output.images[0, -3:, -3:, -1]

								assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()).max() < 1e-4


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




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




					@property
					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Tuple        =     ort.SessionOptions()
								UpperCAmelCase__  :					Dict        =     False
								return options




					def 				snake_case__						(      self):
								# using the PNDM scheduler by default
								UpperCAmelCase__  :					int        =     OnnxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""							,							revision="""onnx"""							,							safety_checker=_lowerCamelCase							,							feature_extractor=_lowerCamelCase							,							provider=self.gpu_provider							,							sess_options=self.gpu_options							,							)
								sd_pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Optional[Any]        =     """A painting of a squirrel eating a burger"""
								np.random.seed(0)
								UpperCAmelCase__  :					str        =     sd_pipe([prompt]							,							guidance_scale=6.0							,							num_inference_steps=10							,							output_type="""np""")
								UpperCAmelCase__  :					Dict        =     output.images

								UpperCAmelCase__  :					Dict        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase__  :					str        =     np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720])

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




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Tuple        =     DDIMScheduler.from_pretrained(
								    """runwayml/stable-diffusion-v1-5"""							,							subfolder="""scheduler"""							,							revision="""onnx""")
								UpperCAmelCase__  :					Optional[int]        =     OnnxStableDiffusionPipeline.from_pretrained(
								    """runwayml/stable-diffusion-v1-5"""							,							revision="""onnx"""							,							scheduler=_lowerCamelCase							,							safety_checker=_lowerCamelCase							,							feature_extractor=_lowerCamelCase							,							provider=self.gpu_provider							,							sess_options=self.gpu_options							,							)
								sd_pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					Tuple        =     """open neural network exchange"""
								UpperCAmelCase__  :					str        =     np.random.RandomState(0)
								UpperCAmelCase__  :					Optional[int]        =     sd_pipe([prompt]							,							guidance_scale=7.5							,							num_inference_steps=10							,							generator=_lowerCamelCase							,							output_type="""np""")
								UpperCAmelCase__  :					List[Any]        =     output.images
								UpperCAmelCase__  :					Optional[Any]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase__  :					int        =     np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486])

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




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[str]        =     LMSDiscreteScheduler.from_pretrained(
								    """runwayml/stable-diffusion-v1-5"""							,							subfolder="""scheduler"""							,							revision="""onnx""")
								UpperCAmelCase__  :					Dict        =     OnnxStableDiffusionPipeline.from_pretrained(
								    """runwayml/stable-diffusion-v1-5"""							,							revision="""onnx"""							,							scheduler=_lowerCamelCase							,							safety_checker=_lowerCamelCase							,							feature_extractor=_lowerCamelCase							,							provider=self.gpu_provider							,							sess_options=self.gpu_options							,							)
								sd_pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					List[str]        =     """open neural network exchange"""
								UpperCAmelCase__  :					Dict        =     np.random.RandomState(0)
								UpperCAmelCase__  :					List[str]        =     sd_pipe([prompt]							,							guidance_scale=7.5							,							num_inference_steps=10							,							generator=_lowerCamelCase							,							output_type="""np""")
								UpperCAmelCase__  :					Any        =     output.images
								UpperCAmelCase__  :					Union[str, Any]        =     image[0, -3:, -3:, -1]

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase__  :					Union[str, Any]        =     np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161])

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




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[Any]        =     0

								def test_callback_fn(_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase) -> None:
											UpperCAmelCase__  :					str        =     True
											nonlocal number_of_steps
											number_of_steps += 1
											if step == 0:
														assert latents.shape == (1, 4, 64, 64)
														UpperCAmelCase__  :					Optional[int]        =     latents[0, -3:, -3:, -1]
														UpperCAmelCase__  :					Any        =     np.array(
														    [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167])

														assert np.abs(latents_slice.flatten() - expected_slice).max() < 1e-3
											elif step == 5:
														assert latents.shape == (1, 4, 64, 64)
														UpperCAmelCase__  :					Tuple        =     latents[0, -3:, -3:, -1]
														UpperCAmelCase__  :					List[str]        =     np.array(
														    [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875])

														assert np.abs(latents_slice.flatten() - expected_slice).max() < 1e-3

								UpperCAmelCase__  :					Optional[Any]        =     False

								UpperCAmelCase__  :					int        =     OnnxStableDiffusionPipeline.from_pretrained(
								    """runwayml/stable-diffusion-v1-5"""							,							revision="""onnx"""							,							safety_checker=_lowerCamelCase							,							feature_extractor=_lowerCamelCase							,							provider=self.gpu_provider							,							sess_options=self.gpu_options							,							)
								pipe.set_progress_bar_config(disable=_lowerCamelCase)

								UpperCAmelCase__  :					List[str]        =     """Andromeda galaxy in a bottle"""

								UpperCAmelCase__  :					int        =     np.random.RandomState(0)
								pipe(
								    prompt=_lowerCamelCase							,							num_inference_steps=5							,							guidance_scale=7.5							,							generator=_lowerCamelCase							,							callback=_lowerCamelCase							,							callback_steps=1							,							)
								assert test_callback_fn.has_been_called
								assert number_of_steps == 6




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Tuple        =     OnnxStableDiffusionPipeline.from_pretrained(
								    """runwayml/stable-diffusion-v1-5"""							,							revision="""onnx"""							,							safety_checker=_lowerCamelCase							,							feature_extractor=_lowerCamelCase							,							provider=self.gpu_provider							,							sess_options=self.gpu_options							,							)
								assert isinstance(_lowerCamelCase							,							_lowerCamelCase)
								assert pipe.safety_checker is None

								UpperCAmelCase__  :					Tuple        =     pipe("""example prompt"""							,							num_inference_steps=2).images[0]
								assert image is not None

								# check that there's no error when saving a pipeline with one of the models being None
								with tempfile.TemporaryDirectory() as tmpdirname:
											pipe.save_pretrained(_lowerCamelCase)
											UpperCAmelCase__  :					Optional[int]        =     OnnxStableDiffusionPipeline.from_pretrained(_lowerCamelCase)

								# sanity check that the pipeline still works
								assert pipe.safety_checker is None
								UpperCAmelCase__  :					Tuple        =     pipe("""example prompt"""							,							num_inference_steps=2).images[0]
								assert image is not None



'''simple docstring'''



from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)


class 		_snake_case					(     a__	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase=None							,							_lowerCamelCase=None							,							_lowerCamelCase=0):
								UpperCAmelCase__  :					Dict        =     1.0 if scale is None else scale
								UpperCAmelCase__  :					Dict        =     0.0 if loc is None else loc

								super().__init__(_lowerCamelCase							,							[AffineTransform(loc=self.loc							,							scale=self.scale							,							event_dim=_lowerCamelCase)])




					@property
					def 				snake_case__						(      self):
								return self.base_dist.mean * self.scale + self.loc




					@property
					def 				snake_case__						(      self):
								return self.base_dist.variance * self.scale**2




					@property
					def 				snake_case__						(      self):
								return self.variance.sqrt()


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							**_lowerCamelCase):
								super().__init__(**_lowerCamelCase)
								UpperCAmelCase__  :					int        =     args_dim
								UpperCAmelCase__  :					Optional[int]        =     nn.ModuleList([nn.Linear(_lowerCamelCase							,							_lowerCamelCase) for dim in args_dim.values()])
								UpperCAmelCase__  :					Optional[Any]        =     domain_map




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					Optional[int]        =     [proj(_lowerCamelCase) for proj in self.proj]

								return self.domain_map(*_lowerCamelCase)


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase):
								super().__init__()
								UpperCAmelCase__  :					Optional[int]        =     function




					def 				snake_case__						(      self							,							_lowerCamelCase							,							*_lowerCamelCase):
								return self.function(_lowerCamelCase							,							*_lowerCamelCase)


class 		_snake_case					:
					lowerCAmelCase      :type
					lowerCAmelCase      :int
					lowerCAmelCase      :Dict[str, int]




					def __init__(      self							,							_lowerCamelCase = 1):
								UpperCAmelCase__  :					Optional[Any]        =     dim
								UpperCAmelCase__  :					int        =     {k: dim * self.args_dim[k] for k in self.args_dim}




					def 				snake_case__						(      self							,							_lowerCamelCase):
								if self.dim == 1:
											return self.distribution_class(*_lowerCamelCase)
								else:
											return Independent(self.distribution_class(*_lowerCamelCase)							,							1)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							):
								UpperCAmelCase__  :					Dict        =     self._base_distribution(_lowerCamelCase)
								if loc is None and scale is None:
											return distr
								else:
											return AffineTransformed(_lowerCamelCase							,							loc=_lowerCamelCase							,							scale=_lowerCamelCase							,							event_dim=self.event_dim)




					@property
					def 				snake_case__						(      self):
								return () if self.dim == 1 else (self.dim,)




					@property
					def 				snake_case__						(      self):
								return len(self.event_shape)




					@property
					def 				snake_case__						(      self):
								return 0.0




					def 				snake_case__						(      self							,							_lowerCamelCase):
								return ParameterProjection(
								    in_features=_lowerCamelCase							,							args_dim=self.args_dim							,							domain_map=LambdaLayer(self.domain_map)							,							)




					def 				snake_case__						(      self							,							*_lowerCamelCase):
								raise NotImplementedError()




					@staticmethod
					def 				snake_case__						(      _lowerCamelCase):
								return (x + torch.sqrt(torch.square(_lowerCamelCase) + 4.0)) / 2.0


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Dict[str, int]   =					{"df": 1, "loc": 1, "scale": 1}
					lowerCAmelCase      :type   =					StudentT




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					Optional[Any]        =     cls.squareplus(_lowerCamelCase).clamp_min(torch.finfo(scale.dtype).eps)
								UpperCAmelCase__  :					Any        =     2.0 + cls.squareplus(_lowerCamelCase)
								return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Dict[str, int]   =					{"loc": 1, "scale": 1}
					lowerCAmelCase      :type   =					Normal




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					str        =     cls.squareplus(_lowerCamelCase).clamp_min(torch.finfo(scale.dtype).eps)
								return loc.squeeze(-1), scale.squeeze(-1)


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Dict[str, int]   =					{"total_count": 1, "logits": 1}
					lowerCAmelCase      :type   =					NegativeBinomial




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					str        =     cls.squareplus(_lowerCamelCase)
								return total_count.squeeze(-1), logits.squeeze(-1)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__   ,       UpperCAmelCase__  :					Tuple        =     distr_args
								if self.dim == 1:
											return self.distribution_class(total_count=_lowerCamelCase							,							logits=_lowerCamelCase)
								else:
											return Independent(self.distribution_class(total_count=_lowerCamelCase							,							logits=_lowerCamelCase)							,							1)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = None):
								UpperCAmelCase__   ,       UpperCAmelCase__  :					Any        =     distr_args

								if scale is not None:
											# See scaling property of Gamma.
											logits += scale.log()

								return self._base_distribution((total_count, logits))