def __call__(
        self, inputs: Array, update_singular_value_estimate: Optional[bool] = None
    ) -> Array:
        """
        Applies a linear transformation with spectral normalization to the inputs.

        Parameters
        ----------
        inputs: Array
            The nd-array to be transformed.
        update_singular_value_estimate: Optional[bool]
            Whether to perform power iterations to update the singular value estimate.

        Returns
        -------
        Array
            The transformed input.
        """
        update_singular_value_estimate = nn.merge_param(
            "update_singular_value_estimate",
            self.update_singular_value_estimate,
            update_singular_value_estimate,
        )
        layer_name = (
            type(self.layer).__name__ if self.layer_name is None else self.layer_name
        )
        params = self.param(
            layer_name, lambda *args: self.layer.init(*args)["params"], inputs
        )
        w = params[self.kernel_name]

        if self.kernel_apply_kwargs is None:
            # By default, we use the implementation in SN-GAN.
            kernel_apply = lambda x: x @ w.reshape(-1, w.shape[-1])
            in_shape = (np.prod(w.shape[:-1]),)
        else:
            # Otherwise, we extract the actual kernel transformation in the input
            # layer. This is useful for Conv2D spectral normalization in
            # [Farzan F. et al., 2019](https://arxiv.org/abs/1811.07457).
            kernel_apply = self.layer.clone(
                **self.kernel_apply_kwargs
            ).bind(  # pylint: disable=not-a-mapping
                {"params": {self.kernel_name: w}}
            )
            # Compute input shape of the kernel operator. This is correct for all
            # linear layers on Flax: Dense, Conv, Embed.
            in_shape = inputs.shape[-w.ndim + 1 : -1] + w.shape[-2:-1]

        initializing = self.is_mutable_collection("params")
        u, v = self._get_singular_vectors(initializing, kernel_apply, in_shape, w.dtype)
        u_hat, v_hat = u.value, v.value
        u_, kernel_transpose = jax.vjp(kernel_apply, v_hat)
        if update_singular_value_estimate and not initializing:
            # Run power iterations using autodiff approach inspired by
            # (Stephan H., 2020)[https://nbviewer.jupyter.org/gist/shoyer/fa9a29fd0880e2e033d7696585978bfc]).
            def scan_body(carry, _):
                u_hat, v_hat, u_ = carry
                (v_,) = kernel_transpose(u_hat)
                v_hat = _l2_normalize(v_)
                u_ = kernel_apply(v_hat)
                u_hat = _l2_normalize(u_)
                return (u_hat, v_hat, u_), None

            (u_hat, v_hat, u_), _ = jax.lax.scan(
                scan_body, (u_hat, v_hat, u_), None, length=self.iteration
            )
            u.value, v.value = u_hat, v_hat

        sigma = jnp.vdot(u_hat, u_)
        # Bound spectral norm by the `norm_multiplier`.
        sigma = jnp.maximum(sigma / self.norm_multiplier, 1.0)
        w_hat = w / jax.lax.stop_gradient(sigma)
        self.sow("intermediates", "w", w_hat)

        # Update params.
        params = flax.core.unfreeze(params)
        params[self.kernel_name] = w_hat
        layer_params = flax.core.freeze({"params": params})
        return self.layer.apply(layer_params, inputs)

        Applies a linear transformation with spectral normalization to the inputs.

        Parameters
        ----------
        inputs: Array
            The nd-array to be transformed.
        update_singular_value_estimate: Optional[bool]
            Whether to perform power iterations to update the singular value estimate.

        Returns
        -------
        Array
            The transformed input.
        
fortuna/model/utils/spectral_norm.py