def init_params(
        self,
        rng: jax.Array,
        mean: Optional[jnp.ndarray] = None,
        log_std: Optional[jnp.ndarray] = None,
    ) -> Dict[str, jnp.ndarray]:
        """
        Initialize architecture parameters.

        Parameters
        ----------
        rng: jax.Array
            Random number generator.
        mean: jnp.ndarray
            If the main model has already been initialized calling `model.init`, the already
            initialized mean parameter values can be provided here.
        log_std: jnp.ndarray
            If the main model has already been initialized calling `model.init`, the already
            initialized log-std parameter values can be provided here.

        Returns
        -------
        Tuple[jnp.array, jnp.array]
            Transformation parameters.
        """
        rng, key_mean, key_log_std = random.split(rng, 3)
        if mean is None:
            mean = self.std_init_params * random.normal(key_mean, (self.dim,))
        if log_std is None:
            log_std = self.std_init_params * random.normal(key_log_std, (self.dim,))
        return dict(mean=mean, log_std=log_std)

        Initialize architecture parameters.

        Parameters
        ----------
        rng: jax.Array
            Random number generator.
        mean: jnp.ndarray
            If the main model has already been initialized calling `model.init`, the already
            initialized mean parameter values can be provided here.
        log_std: jnp.ndarray
            If the main model has already been initialized calling `model.init`, the already
            initialized log-std parameter values can be provided here.

        Returns
        -------
        Tuple[jnp.array, jnp.array]
            Transformation parameters.
        
fortuna/prob_model/posterior/normalizing_flow/advi/advi_architecture.py