code: def make_fast_generalized_attention(
      qkv_dim,
      renormalize_attention=True,
      numerical_stabilizer=0.0,
      nb_features=256,
      features_type="deterministic",
      kernel_fn=jax.nn.relu,
      kernel_epsilon=0.001,
      redraw_features=False,
      unidirectional=False,
      lax_scan_unroll=1,
  ):
      """Construct a fast generalized attention menthod."""
      logging.info(
          "Fast generalized attention.: %s features and renormalize=%s",
          nb_features,
          renormalize_attention,
      )
      if features_type == "ortho":
          matrix_creator = functools.partial(
              GaussianOrthogonalRandomMatrix,
              nb_features,
              qkv_dim,
              scaling=False,
          )
      elif features_type == "iid":
          matrix_creator = functools.partial(
              GaussianUnstructuredRandomMatrix, nb_features, qkv_dim
          )
      elif features_type == "deterministic":
          matrix_creator = None
      else:
          raise ValueError("Unknown feature value type")
  
      def kernel_feature_creator(
          data,
          projection_matrix,
          attention_dims_t,
          batch_dims_t,
          precision,
          is_query,
          normalize_data=False,
      ):
          del attention_dims_t
          del is_query
          return generalized_kernel_feature_creator(
              data,
              projection_matrix,
              batch_dims_t,
              precision,
              kernel_fn,
              kernel_epsilon,
              normalize_data,
          )
  
      attention_fn = FastAttentionviaLowRankDecomposition(
          matrix_creator,
          kernel_feature_creator,
          renormalize_attention=renormalize_attention,
          numerical_stabilizer=numerical_stabilizer,
          redraw_features=redraw_features,
          unidirectional=unidirectional,
          lax_scan_unroll=lax_scan_unroll,
      ).dot_product_attention
      return attention_fn
  
imports: import abc
  import functools
  from collections.abc import Iterable
  import numpy
  from absl import logging
  import jax
  import jax.numpy
  from jax import lax, random
label_desc: Use image node_id 5 for calling a global function with example usage: make_fast_generalized_attention(qkv_dim, renormalize_attention, numerical_stabilizer, nb_features, features_type, kernel_fn, kernel_epsilon, redraw_features, unidirectional, lax_scan_unroll) and returns: attention_fn, generalized_kernel_feature_creator