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

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

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

        Component-wise backward linear transformation.

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

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