



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




		'''simple docstring'''



		def __init__(self	,__lowerCamelCase	,__lowerCamelCase=None	,__lowerCamelCase=None	,__lowerCamelCase=0				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  1.0 if scale is None else scale
						lowerCAmelCase__  :       Optional[int]				=  0.0 if loc is None else loc

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



		@property
		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						return self.base_dist.mean * self.scale + self.loc



		@property
		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






						return self.base_dist.variance * self.scale**2



		@property
		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						return self.variance.sqrt()



class  lowerCamelCase__    (							nn.Module):




		'''simple docstring'''



		def __init__(self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,**__lowerCamelCase				)  ->    None:

						"""simple docstring"""






						super().__init__(**__lowerCamelCase				)
						lowerCAmelCase__  :       Tuple				=  args_dim
						lowerCAmelCase__  :       Dict				=  nn.ModuleList([nn.Linear(__lowerCamelCase	,__lowerCamelCase				) for dim in args_dim.values()]				)
						lowerCAmelCase__  :       List[str]				=  domain_map



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Tuple[torch.Tensor]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  [proj(__lowerCamelCase				) for proj in self.proj]

						return self.domain_map(*__lowerCamelCase				)



class  lowerCamelCase__    (							nn.Module):




		'''simple docstring'''



		def __init__(self	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						super().__init__()
						lowerCAmelCase__  :       Optional[int]				=  function



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,*__lowerCamelCase				)  ->    Union[str, Any]:

						"""simple docstring"""






						return self.function(__lowerCamelCase	,*__lowerCamelCase				)



class  lowerCamelCase__    :




		'''simple docstring'''



		snake_case_			=42
		snake_case_			=42
		snake_case_			=42



		def __init__(self	,__lowerCamelCase = 1				)  ->    None:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  dim
						lowerCAmelCase__  :       int				=  {k: dim * self.args_dim[k] for k in self.args_dim}



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						if self.dim == 1:
										return self.distribution_class(*__lowerCamelCase				)
						else:
										return Independent(self.distribution_class(*__lowerCamelCase				)	,1				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase = None	,__lowerCamelCase = None	,)  ->    Distribution:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  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 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						return () if self.dim == 1 else (self.dim,)



		@property
		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						return len(self.event_shape				)



		@property
		def 	lowerCAmelCase__   (self				)  ->    float:

						"""simple docstring"""






						return 0.0



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    nn.Module:

						"""simple docstring"""






						return ParameterProjection(
						    in_features=__lowerCamelCase	,args_dim=self.args_dim	,domain_map=LambdaLayer(self.domain_map				)	,)



		def 	lowerCAmelCase__   (self	,*__lowerCamelCase				)  ->    Tuple:

						"""simple docstring"""






						raise NotImplementedError()



		@staticmethod
		def 	lowerCAmelCase__   (__lowerCamelCase				)  ->    torch.Tensor:

						"""simple docstring"""






						return (x + torch.sqrt(torch.square(__lowerCamelCase				) + 4.0				)) / 2.0



class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_			={"df": 1, "loc": 1, "scale": 1}
		snake_case_			=StudentT



		@classmethod
		def 	lowerCAmelCase__   (cls	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  cls.squareplus(__lowerCamelCase				).clamp_min(torch.finfo(scale.dtype				).eps				)
						lowerCAmelCase__  :       str				=  2.0 + cls.squareplus(__lowerCamelCase				)
						return df.squeeze(-1				), loc.squeeze(-1				), scale.squeeze(-1				)



class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_			={"loc": 1, "scale": 1}
		snake_case_			=Normal



		@classmethod
		def 	lowerCAmelCase__   (cls	,__lowerCamelCase	,__lowerCamelCase				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  cls.squareplus(__lowerCamelCase				).clamp_min(torch.finfo(scale.dtype				).eps				)
						return loc.squeeze(-1				), scale.squeeze(-1				)



class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_			={"total_count": 1, "logits": 1}
		snake_case_			=NegativeBinomial



		@classmethod
		def 	lowerCAmelCase__   (cls	,__lowerCamelCase	,__lowerCamelCase				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  cls.squareplus(__lowerCamelCase				)
						return total_count.squeeze(-1				), logits.squeeze(-1				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Distribution:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       Optional[Any]				=  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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase = None	,__lowerCamelCase = None				)  ->    Distribution:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       Tuple				=  distr_args

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

						return self._base_distribution((total_count, logits)				)






from collections.abc import Callable

import numpy as np






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Callable ,lowerCamelCase_	:	float ,lowerCamelCase_	:	float ,lowerCamelCase_	:	float ,lowerCamelCase_	:	float):





				'''simple docstring'''

				lowerCAmelCase__  :       Dict				=  int(np.ceil((x_end - xa) / step_size))
				lowerCAmelCase__  :       Any				=  np.zeros((n + 1,))
				lowerCAmelCase__  :       Union[str, Any]				=  ya
				lowerCAmelCase__  :       str				=  xa

				for k in range(lowerCamelCase_):
								lowerCAmelCase__  :       Tuple				=  y[k] + step_size * ode_func(lowerCamelCase_ ,y[k])
								lowerCAmelCase__  :       str				=  y[k] + (
								    (step_size / 2) * (ode_func(lowerCamelCase_ ,y[k]) + ode_func(x + step_size ,lowerCamelCase_))
								)
								x += step_size

				return y


if __name__ == "__main__":
							import doctest

							doctest.testmod()

