docstr: Unified Skew Normal distribution of `Y | a < X < b` for jointly Gaussian
  random vectors `X ∈ R^m` and `Y ∈ R^n`.
  
  Batch shapes `trunc.batch_shape` and `gauss.batch_shape` must be broadcastable.
  Care should be taken when choosing `trunc.batch_shape`. When `trunc` is of lower
  batch dimensionality than `gauss`, the user should consider expanding `trunc` to
  hasten `UnifiedSkewNormal.log_prob`. In these cases, it is suggested that the
  user invoke `trunc.solver` before calling `trunc.expand` to avoid paying for
  multiple, identical solves.
  
  Args:
      trunc: Distribution of `Z = (X | a < X < b) ∈ R^m`.
      gauss: Distribution of `Y ∈ R^n`.
      cross_covariance_matrix: Cross-covariance `Cov(X, Y) ∈ R^{m x n}`.
      validate_args: Optional argument to super().__init__.
docstr_tok: ["Unified","Skew","Normal","distribution","of","`","Y","|","a","<","X","<","b","`","for","jointly","Gaussian","random","vectors","`","X","\u2208","R^m","`","and","`","Y","\u2208","R^n","`",".","Batch","shapes","`","trunc.batch_shape","`","and","`","gauss.batch_shape","`","must","be","broadcastable",".","Care","should","be","taken","when","choosing","`","trunc.batch_shape","`",".","When","`","trunc","`","is","of","lower","batch","dimensionality","than","`","gauss","`",",","the","user","should","consider","expanding","`","trunc","`","to","hasten","`","UnifiedSkewNormal.log_prob","`",".","In","these","cases",",","it","is","suggested","that","the","user","invoke","`","trunc.solver","`","before","calling","`","trunc.expand","`","to","avoid","paying","for","multiple",",","identical","solves",".","Args",":","trunc",":","Distribution","of","`","Z","=","(","X","|","a","<","X","<","b",")","\u2208","R^m","`",".","gauss",":","Distribution","of","`","Y","\u2208","R^n","`",".","cross_covariance_matrix",":","Cross-covariance","`","Cov","(","X",",","Y",")","\u2208","R^","{","m","x","n","}","`",".","validate_args",":","Optional","argument","to","super","(",")",".__init__","."]
code: def __init__(
      self,
      trunc: TruncatedMultivariateNormal,
      gauss: MultivariateNormal,
      cross_covariance_matrix: Union[Tensor, LinearOperator],
      validate_args: Optional[bool] = None,
  ):
      r"""Unified Skew Normal distribution of `Y | a < X < b` for jointly Gaussian
      random vectors `X ∈ R^m` and `Y ∈ R^n`.
  
      Batch shapes `trunc.batch_shape` and `gauss.batch_shape` must be broadcastable.
      Care should be taken when choosing `trunc.batch_shape`. When `trunc` is of lower
      batch dimensionality than `gauss`, the user should consider expanding `trunc` to
      hasten `UnifiedSkewNormal.log_prob`. In these cases, it is suggested that the
      user invoke `trunc.solver` before calling `trunc.expand` to avoid paying for
      multiple, identical solves.
  
      Args:
          trunc: Distribution of `Z = (X | a < X < b) ∈ R^m`.
          gauss: Distribution of `Y ∈ R^n`.
          cross_covariance_matrix: Cross-covariance `Cov(X, Y) ∈ R^{m x n}`.
          validate_args: Optional argument to super().__init__.
      """
      if len(trunc.event_shape) != len(gauss.event_shape):
          raise ValueError(
              f"{len(trunc.event_shape)}-dimensional `trunc` incompatible with"
              f"{len(gauss.event_shape)}-dimensional `gauss`."
          )
      # LinearOperator currently doesn't support torch.linalg.solve_triangular,
      # so for the time being, we cast the operator to dense here
      if isinstance(cross_covariance_matrix, LinearOperator):
          cross_covariance_matrix = cross_covariance_matrix.to_dense()
      try:
          batch_shape = torch.broadcast_shapes(
              trunc.batch_shape, gauss.batch_shape
          )
      except RuntimeError as e:
          raise ValueError("Incompatible batch shapes") from e
  
      super().__init__(
          batch_shape=batch_shape,
          event_shape=gauss.event_shape,
          validate_args=validate_args,
      )
      self.trunc = trunc
      self.gauss = gauss
      self.cross_covariance_matrix = cross_covariance_matrix
      if self._validate_args:
          try:
              # calling _orthogonalized_gauss first makes the following call
              # _orthogonalized_gauss.scale_tril which is used by self.rsample
              self._orthogonalized_gauss
              self.scale_tril
          except Exception as e:
              # error could be thrown by linalg.augment_cholesky (NotPSDError)
              # or torch.linalg.cholesky (with "positive-definite" in the message)
              if (
                  isinstance(e, NotPSDError)
                  or "positive-definite" in str(e)
                  or "PositiveDefinite" in str(e)
              ):
                  e = ValueError(
                      "UnifiedSkewNormal is only well-defined for positive definite"
                      " joint covariance matrices."
                  )
              raise e
  
imports: from __future__ import annotations
  from inspect import getmembers
  from typing import Optional, Sequence, Union
  import torch
  from botorch.utils.probability.linalg import augment_cholesky, block_matrix_concat
  from botorch.utils.probability.mvnxpb import MVNXPB
  from botorch.utils.probability.truncated_multivariate_normal import TruncatedMultivariateNormal
  from linear_operator.operators import LinearOperator
  from linear_operator.utils.errors import NotPSDError
  from torch import Tensor
  from torch.distributions.multivariate_normal import Distribution, MultivariateNormal
  from torch.distributions.utils import lazy_property
  from torch.nn.functional import pad
label_desc: Use image node_id 1 to create a new UnifiedSkewNormal object from inherited base classes: Distribution with example: obj = UnifiedSkewNormal(trunc, gauss, cross_covariance_matrix, validate_args)