def __call__(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray:
        """
        Forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Outputs of the hidden layers.

        Returns
        -------
        jnp.ndarray
            Output of the last layer.
        """
        if self.activation is not None:
            x = self.activation(x)
        x = self.dense(self.output_dim, name="last")(x)
        return x

        Forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Outputs of the hidden layers.

        Returns
        -------
        jnp.ndarray
            Output of the last layer.
        
fortuna/model/mlp.py