code: def sincos_softmax_kernel_feature_creator(
      data,
      projection_matrix,
      attention_dims_t,
      batch_dims_t,
      precision,
      normalize_data=True,
  ):
      """Constructs kernel sin-cos features for fast softmax attention.
  
      Args:
        data: input for which features are computes
        projection_matrix: random matrix used to compute features
        attention_dims_t: tuple of attention dimensions
        batch_dims_t: tuple of batch dimensions
        precision: precision parameter
        normalize_data: predicate indicating whether data should be normalized.
  
      Returns:
        Random features for fast softmax attention.
      """
      if normalize_data:
          # We have: exp(qk^T/sqrt{d}) = exp(|q|^2/2sqrt{d}) * exp(|k|^2/2sqrt{d}) *
          # exp(-(|q*c-k*c|^2)/2), where c = 1.0 / sqrt{sqrt{d}}.
          data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
      else:
          data_normalizer = 1.0
      ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0])
      data_mod_shape = (
          data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
      )
      data_thick_random_matrix = (
          jnp.zeros(data_mod_shape) + projection_matrix
      )
  
      data_dash = lax.dot_general(
          data_normalizer * data,
          data_thick_random_matrix,
          (
              ((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)),
              (batch_dims_t, batch_dims_t),
          ),
          precision=precision,
      )
      data_dash_cos = ratio * jnp.cos(data_dash)
      data_dash_sin = ratio * jnp.sin(data_dash)
      data_dash = jnp.concatenate(
          (data_dash_cos, data_dash_sin), axis=-1
      )
  
      # Constructing D_data and data^{'}
      diag_data = jnp.square(data)
      diag_data = jnp.sum(diag_data, axis=data.ndim - 1)
      diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer
      diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1)
      # Additional renormalization for numerical stability
      data_renormalizer = jnp.max(
          diag_data, attention_dims_t, keepdims=True
      )
      diag_data -= data_renormalizer
      diag_data = jnp.exp(diag_data)
      data_prime = data_dash * diag_data
      return data_prime
  
imports: import abc
  import enum
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Iterable
  from absl import logging
  from flax import linen
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 7 for calling a global function with example usage: sincos_softmax_kernel_feature_creator(data, projection_matrix, attention_dims_t, batch_dims_t, precision, normalize_data) and returns: data_prime