def init(self, input_shape: Shape, **kwargs) -> JointState:
        """
        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.likelihood.model_manager.init(
                input_shape, rng=self.rng.get(), **kwargs
            )
        )
        inputs = get_inputs_from_shape(input_shape)
        outputs = self.likelihood.model_manager.apply(
            params=oms.params, inputs=inputs, mutable=oms.mutable
        )
        output_dim = (
            outputs[0].shape[-1]
            if isinstance(outputs, (list, tuple))
            else outputs.shape[-1]
        )
        ocms = OutputCalibManagerState.init_from_dict(
            FrozenDict(
                output_calibrator=self.likelihood.output_calib_manager.init(
                    output_dim=output_dim
                )
            )
        )
        return JointState.init_from_states(oms, ocms)

        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/prob_model/joint/base.py