def _get_output_dim(self, input_shape: Shape, **kwargs) -> int:
        """
        Initialize the state of the joint distribution.

        Parameters
        ----------
        input_shape : Shape
            The shape of the input variable.

        Returns
        -------
        A state of the joint distribution.
        """
        oms = ModelManagerState.init_from_dict(
            self.model_manager.init(input_shape, rng=self.rng.get(), **kwargs)
        )
        inputs = get_inputs_from_shape(input_shape)
        outputs = self.model_manager.apply(oms.params, inputs, mutable=oms.mutable)
        return (
            outputs[0].shape[-1]
            if isinstance(outputs, (list, tuple))
            else outputs.shape[-1]
        )

        Initialize the state of the joint distribution.

        Parameters
        ----------
        input_shape : Shape
            The shape of the input variable.

        Returns
        -------
        A state of the joint distribution.
        
fortuna/calib_model/base.py