code: def make_fast_softmax_attention(
      qkv_dim,
      renormalize_attention=True,
      numerical_stabilizer=0.000001,
      nb_features=256,
      ortho_features=True,
      ortho_scaling=0.0,
      redraw_features=True,
      unidirectional=False,
      nonnegative_features=True,
      lax_scan_unroll=1,
  ):
      """Construct a fast softmax attention method."""
      """
    logging.info(
        'Fast softmax attention: %s features and orthogonal=%s, renormalize=%s',
        nb_features, ortho_features, renormalize_attention)
    """
      if ortho_features:
          matrix_creator = functools.partial(
              GaussianOrthogonalRandomMatrix,
              nb_features,
              qkv_dim,
              scaling=ortho_scaling,
          )
      else:
          matrix_creator = functools.partial(
              GaussianUnstructuredRandomMatrix, nb_features, qkv_dim
          )
      if nonnegative_features:
  
          def kernel_feature_creator(
              data,
              projection_matrix,
              attention_dims_t,
              batch_dims_t,
              precision,
              is_query,
              normalize_data=True,
          ):
              return nonnegative_softmax_kernel_feature_creator(
                  data,
                  projection_matrix,
                  attention_dims_t,
                  batch_dims_t,
                  precision,
                  is_query,
                  normalize_data,
                  numerical_stabilizer,
              )
  
      else:
  
          def kernel_feature_creator(
              data,
              projection_matrix,
              attention_dims_t,
              batch_dims_t,
              precision,
              is_query,
              normalize_data=True,
          ):
              del is_query
              return sincos_softmax_kernel_feature_creator(
                  data,
                  projection_matrix,
                  attention_dims_t,
                  batch_dims_t,
                  precision,
                  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 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 9 for calling a global function with example usage: make_fast_softmax_attention(qkv_dim, renormalize_attention, numerical_stabilizer, nb_features, ortho_features, ortho_scaling, redraw_features, unidirectional, nonnegative_features, lax_scan_unroll) and returns: attention_fn, nonnegative_softmax_kernel_feature_creator, sincos_softmax_kernel_feature_creator