def __call__(self, inputs: Array) -> Array:
        """
        Applies random feature transformation along the last dimension of inputs.

        Parameters
        ----------
        inputs: Array
            The nd-array to be transformed.

        Returns
        -------
        Array
          The transformed input.
        """
        # Initializes variables.
        input_dim = inputs.shape[-1]

        kernel_rng, bias_rng = random.split(self.rng, num=2)
        kernel_shape = (input_dim, self.features)

        kernel = self.variable(
            self.collection_name,
            "kernel",
            self.kernel_init,
            kernel_rng,
            kernel_shape,
            self.dtype,
        )
        kernel_scale = self.param(
            "rf_kernel_scale",
            nn.initializers.constant(self.kernel_scale),
            (1,),
            self.dtype,
        )
        bias = self.variable(
            self.collection_name,
            "bias",
            self.bias_init,
            bias_rng,
            (self.features,),
            self.dtype,
        )

        # Specifies multiplication dimension.
        contracting_dims = ((inputs.ndim - 1,), (0,))
        batch_dims = ((), ())

        # Performs forward pass.
        inputs = jnp.asarray(inputs, self.dtype)

        outputs = lax.dot_general(
            inputs, (1.0 / kernel_scale) * kernel.value, (contracting_dims, batch_dims)
        )
        outputs = outputs + jnp.broadcast_to(bias.value, outputs.shape)

        return self._feature_scale * jnp.cos(outputs)

        Applies random feature transformation along the last dimension of inputs.

        Parameters
        ----------
        inputs: Array
            The nd-array to be transformed.

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