def forward(
        self, params: Dict[str, jnp.ndarray], u: jnp.ndarray
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Component-wise forward linear transformation.

        Parameters
        ----------
        params: Dict[str, jnp.ndarray]
            Mean and log-std parameters.
        u: jnp.ndarray
            Input to transform forward.

        Returns
        -------
        Tuple[jnp.ndarray, jnp.ndarray]
            v: jnp.ndarray
                Output of the forward pass.
            ldj: jnp.ndarray
                Log-determinant of the Jacobian of the forward pass.
        """
        return (
            params["mean"] + jnp.exp(params["log_std"]) * u,
            jnp.repeat(jnp.sum(params["log_std"], -1), u.shape[0]),
        )

        Component-wise forward linear transformation.

        Parameters
        ----------
        params: Dict[str, jnp.ndarray]
            Mean and log-std parameters.
        u: jnp.ndarray
            Input to transform forward.

        Returns
        -------
        Tuple[jnp.ndarray, jnp.ndarray]
            v: jnp.ndarray
                Output of the forward pass.
            ldj: jnp.ndarray
                Log-determinant of the Jacobian of the forward pass.
        
fortuna/prob_model/posterior/normalizing_flow/advi/advi_architecture.py