def _batched_log_joint_prob(
        self,
        params: Params,
        batch: Batch,
        n_data: int,
        mutable: Optional[Mutable] = None,
        calib_params: Optional[CalibParams] = None,
        calib_mutable: Optional[CalibMutable] = None,
        return_aux: Optional[List[str]] = None,
        train: Optional[bool] = False,
        outputs: Optional[jnp.ndarray] = None,
        rng: Optional[jax.Array] = None,
        **kwargs,
    ) -> Union[float, Tuple[float, dict]]:
        """
        Evaluate the joint batched log-probability density function (a.k.a. log-pdf).

        Parameters
        ----------
        params : Params
            The random parameters of the probabilistic model.
        batch : Batch
            A batch of data points.
        n_data : int
            The total number of data points over which the likelihood is joint. This is used to rescale the batched
            log-likelihood function to better approximate the full likelihood.
        mutable : Optional[Mutable]
            The mutable objects used to evaluate the models.
        calib_params : Optional[CalibParams]
            The calibration parameters of the probabilistic model.
        calib_mutable : Optional[CalibMutable]
            The calibration mutable objects used to evaluate the calibrators.
        return_aux : Optional[List[str]]
            The auxiliary objects to return. We support 'outputs' and 'calib_mutable'. If this argument is not given,
            no auxiliary object is returned.
        train : Optional[bool]
            Whether the method is called during training.
        outputs : Optional[jnp.ndarray]
            Pre-computed batch of outputs.
        rng: Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.

        Returns
        -------
        Union[float, Tuple[float, dict]]
            The evaluation of the joint batched log-pdf. If `return_aux` is given, the corresponding auxiliary objects
            are also returned.
        """
        if return_aux is None:
            return_aux = []
        log_prior = self.prior.log_joint_prob(params)
        outs = self.likelihood._batched_log_joint_prob(
            params,
            batch,
            n_data,
            mutable=mutable,
            calib_params=calib_params,
            calib_mutable=calib_mutable,
            return_aux=return_aux,
            train=train,
            outputs=outputs,
            rng=rng,
            **kwargs,
        )
        if len(return_aux) > 0:
            batched_log_lik, aux = outs
            return batched_log_lik + log_prior, aux
        batched_log_lik = outs
        return batched_log_lik + log_prior

        Evaluate the joint batched log-probability density function (a.k.a. log-pdf).

        Parameters
        ----------
        params : Params
            The random parameters of the probabilistic model.
        batch : Batch
            A batch of data points.
        n_data : int
            The total number of data points over which the likelihood is joint. This is used to rescale the batched
            log-likelihood function to better approximate the full likelihood.
        mutable : Optional[Mutable]
            The mutable objects used to evaluate the models.
        calib_params : Optional[CalibParams]
            The calibration parameters of the probabilistic model.
        calib_mutable : Optional[CalibMutable]
            The calibration mutable objects used to evaluate the calibrators.
        return_aux : Optional[List[str]]
            The auxiliary objects to return. We support 'outputs' and 'calib_mutable'. If this argument is not given,
            no auxiliary object is returned.
        train : Optional[bool]
            Whether the method is called during training.
        outputs : Optional[jnp.ndarray]
            Pre-computed batch of outputs.
        rng: Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.

        Returns
        -------
        Union[float, Tuple[float, dict]]
            The evaluation of the joint batched log-pdf. If `return_aux` is given, the corresponding auxiliary objects
            are also returned.
        
fortuna/prob_model/joint/base.py