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 functools
  from collections.abc import Iterable
  import jax
  import jax.numpy
  import numpy
  from absl import logging
  from jax import lax, random
label_desc: Use image node_id 4 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