def __call__(self, x: Array, train: bool = False, **kwargs) -> jnp.ndarray:
        """
        Forward pass.

        Parameters
        ----------
        x: Array
            Inputs.
        train: bool
            Whether it is training or inference.

        Returns
        -------
        jnp.ndarray
            Output of the hidden layers.
        """
        if hasattr(self, "spectral_norm"):
            dense = self.spectral_norm(self.dense, train=train)
        else:
            dense = self.dense
        dropout = self.dropout(self.dropout_rate)
        n_activations = len(self.activations)

        def update(i: int, x):
            x = dense(self.widths[i], name="hidden" + str(i + 1))(x)
            if i < n_activations:
                x = self.activations[i](x)
            x = dropout(x, deterministic=not train)
            return x

        x = x.reshape(x.shape[0], -1)
        for i in range(0, len(self.widths)):
            x = update(i, x)
        return x

        Forward pass.

        Parameters
        ----------
        x: Array
            Inputs.
        train: bool
            Whether it is training or inference.

        Returns
        -------
        jnp.ndarray
            Output of the hidden layers.
        
fortuna/model/mlp.py