tj:,ddlZddlZddlmZddlmZddlmZddlmZm Z ddl m Z dgZ dZ d Zd ZGd deZdS) N)Tensor) constraints) Distribution)_standard_normal lazy_property)_sizeMultivariateNormalcxtj||ddS)a Performs a batched matrix-vector product, with compatible but different batch shapes. This function takes as input `bmat`, containing :math:`n \times n` matrices, and `bvec`, containing length :math:`n` vectors. Both `bmat` and `bvec` may have any number of leading dimensions, which correspond to a batch shape. They are not necessarily assumed to have the same batch shape, just ones which can be broadcasted. )torchmatmul unsqueezesqueeze)bmatbvecs e/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/torch/distributions/multivariate_normal.py _batch_mvrs0 <dnnR00 1 1 9 9" = ==c|d}|jdd}t|}|dz }||z }||z}|d|zz}|jd|} t |jdd|j|dD]\} } | | | z| fz } | |fz } || }t t|t t||dzt t|dz|dz|gz} || }|d||} |d| d|}|ddd}tj | |d d d}|}||jdd}t t|}t|D]}|||z||zgz }||}||S) aN Computes the squared Mahalanobis distance :math:`\mathbf{x}^\top\mathbf{M}^{-1}\mathbf{x}` for a factored :math:`\mathbf{M} = \mathbf{L}\mathbf{L}^\top`. Accepts batches for both bL and bx. They are not necessarily assumed to have the same batch shape, but `bL` one should be able to be broadcasted to `bx` one. r NrFupper)sizeshapelendimzipreshapelistrangepermuter linalgsolve_triangularpowsumt)bLbxnbx_batch_shape bx_batch_dims bL_batch_dimsouter_batch_dimsold_batch_dimsnew_batch_dims bx_new_shapesLsx permute_dimsflat_Lflat_x flat_x_swapM_swapM permuted_Mpermute_inv_dimsi reshaped_Ms r_batch_mahalanobisr?s  AXcrc]N''MFFHHqLM$}4% 5N%M(99N8---.LbhssmRX.>r.A%BCC''Br2& QDL L ! !B U# $ $%% u%~q99 : : ; u%)>1== > > ?    L ! !B ZZAq ! !F ZZFKKNNA . .F..Aq))K %%fk%GGKKANNRRSUVV   A28CRC=))JE"23344 = ! !GG-1>A3EFF##$455J   n - --rcXtjtj|d}tjtj|ddd}tj|jd|j|j}tj ||d}|S)N)rr rr dtypedeviceFr) r r$choleskyflip transposeeyerrBrCr%)PLfL_invIdLs r_precision_to_scale_trilrMOs   uz!X66 7 7B OEJr844b" = =E 172;agah ? ? ?B %%eRu%==A Hrc eZdZdZejejejejdZejZ dZ dde de dzde dzde dzd e dzd df fd Z dfd Zed e fd Zed e fdZed e fdZed e fdZed e fdZed e fdZejfded e fdZdZdZxZS)r a Creates a multivariate normal (also called Gaussian) distribution parameterized by a mean vector and a covariance matrix. The multivariate normal distribution can be parameterized either in terms of a positive definite covariance matrix :math:`\mathbf{\Sigma}` or a positive definite precision matrix :math:`\mathbf{\Sigma}^{-1}` or a lower-triangular matrix :math:`\mathbf{L}` with positive-valued diagonal entries, such that :math:`\mathbf{\Sigma} = \mathbf{L}\mathbf{L}^\top`. This triangular matrix can be obtained via e.g. Cholesky decomposition of the covariance. Example: >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) >>> # xdoctest: +IGNORE_WANT("non-deterministic") >>> m = MultivariateNormal(torch.zeros(2), torch.eye(2)) >>> m.sample() # normally distributed with mean=`[0,0]` and covariance_matrix=`I` tensor([-0.2102, -0.5429]) Args: loc (Tensor): mean of the distribution covariance_matrix (Tensor): positive-definite covariance matrix precision_matrix (Tensor): positive-definite precision matrix scale_tril (Tensor): lower-triangular factor of covariance, with positive-valued diagonal Note: Only one of :attr:`covariance_matrix` or :attr:`precision_matrix` or :attr:`scale_tril` can be specified. Using :attr:`scale_tril` will be more efficient: all computations internally are based on :attr:`scale_tril`. If :attr:`covariance_matrix` or :attr:`precision_matrix` is passed instead, it is only used to compute the corresponding lower triangular matrices using a Cholesky decomposition. )loccovariance_matrixprecision_matrix scale_trilTNrOrPrQrR validate_argsreturnc|dkrtd|du|duz|duzdkrtd|t|dkrtdtj|jdd|jdd}||dz|_n|t|dkrtd tj|jdd|jdd}||dz|_n|td |dkrtd tj|jdd|jdd}||dz|_ ||d z|_ |j jdd}t ||| | ||_ dS|&tj||_ dSt!||_ dS)Nrz%loc must be at least one-dimensional.zTExactly one of covariance_matrix or precision_matrix or scale_tril may be specified.rzZscale_tril matrix must be at least two-dimensional, with optional leading batch dimensionsrr )r r zZcovariance_matrix must be at least two-dimensional, with optional leading batch dimensionsz%precision_matrix is unexpectedly NonezYprecision_matrix must be at least two-dimensional, with optional leading batch dimensions)r rS)r ValueErrorr broadcast_shapesrexpandrRrPAssertionErrorrQrOsuper__init___unbroadcasted_scale_trilr$rDrM) selfrOrPrQrRrS batch_shape event_shape __class__s rr\zMultivariateNormal.__init__s 7799q==DEE E T )j.D E D (   f   !~~!## = 01A#2#1F RUSURUWWK(// h0FGGDOO  * $$&&** = 0!',cinK&7%=%=kH>T%U%UD " "'$%LMMM##%%)) = 0 &ss+SYss^K%5$;$;K(>CN  . .#'#8#?#? #J#JC   #&&// ) 0   "0 rc`|j|j|jz|jzSN)r]rY _batch_shape _event_shaper^s rrRzMultivariateNormal.scale_trils3-44   1 1D4E E   rctj|j|jj|j|jz|jzSrl)r r r]mTrYrmrnros rrPz$MultivariateNormal.covariance_matrixsE|  *D,J,M  &"T%669JJ K K Lrctj|j|j|jz|jzSrl)r cholesky_inverser]rYrmrnros rrQz#MultivariateNormal.precision_matrixs>%d&DEELL   1 1D4E E   rc|jSrlrOros rmeanzMultivariateNormal.mean xrc|jSrlruros rmodezMultivariateNormal.moderwrc|jdd|j|jzS)Nrr )r]r&r'rYrmrnros rvariancezMultivariateNormal.variances@  * . .q 1 1 SWW VD%(99 : : r sample_shapec||}t||jj|jj}|jt |j|zS)NrA)_extended_shaperrOrBrCrr])r^r|repss rrsamplezMultivariateNormal.rsamplesK$$\22uDHN48?SSSx)D$BCHHHHrcj|jr||||jz }t|j|}|jddd}d|jdtjdtj zz|zz|z S)Nrr dim1dim2grr) rf_validate_samplerOr?r]diagonallogr'rnmathpi)r^valuediffr: half_log_dets rlog_probzMultivariateNormal.log_probs   )  ! !% ( ( (tx t=t D D  * 3 3" 3 E E I I K K O OPR S S t(+dhq47{.C.CCaGH<WWrc\|jddd}d|jdzdt jdt jzzz|z}t|jdkr|S| |jS)Nrr rg?rg?r) r]rrr'rnrrrrmrY)r^rHs rentropyzMultivariateNormal.entropy s  * 3 3" 3 E E I I K K O OPR S S  $#A& &#TW0E0E*E F U t ! !Q & &H88D-.. .r)NNNNrl)__name__ __module__ __qualname____doc__r real_vectorpositive_definitelower_choleskyarg_constraintssupport has_rsamplerboolr\rYrrRrPrQpropertyrvryr{r rdrrrr __classcell__)ras@rr r XsQ""L&(:'9!0 O %GK ,0*.$(%) =X=X =X"D==X!4- =X TM =X d{ =X =X=X=X=X=X=X~& F   ] L6LLL]L  &   ] fXfX &   X -7EJLLIIEIVIIII XXX///////r)rr rtorch.distributionsr torch.distributions.distributionrtorch.distributions.utilsrr torch.typesr__all__rr?rMr rrrs ++++++999999EEEEEEEE  > > >/././.d   z/z/z/z/z/z/z/z/z/z/r