class_docstr: Self Attention Layer.
docstr: Applies self attention on the input data.
  
  Args:
    inputs: Input tensor of shape [batch_size, num_points, feature_dim]
    coords: Input tensor of point positions of the shape [batch_size,
      num_points, 3]
    mask: Binary array of shape broadcastable to `inputs` tensor, indicating
      the positions for which self attention should be computed.
    numerical_stabilizer: takes into account stability when inputs are zero.
    train: Whether it is training or not.
  
  Returns:
    output: Tensor of shape [batch_size, num_points, feature_dim]
docstr_tok: ["Applies","self","attention","on","the","input","data",".","Args",":","inputs",":","Input","tensor","of","shape","[","batch_size",",","num_points",",","feature_dim","]","coords",":","Input","tensor","of","point","positions","of","the","shape","[","batch_size",",","num_points",",","3","]","mask",":","Binary","array","of","shape","broadcastable","to","`","inputs","`","tensor",",","indicating","the","positions","for","which","self","attention","should","be","computed",".","numerical_stabilizer",":","takes","into","account","stability","when","inputs","are","zero",".","train",":","Whether","it","is","training","or","not",".","Returns",":","output",":","Tensor","of","shape","[","batch_size",",","num_points",",","feature_dim","]"]
code: def __call__(
      self,
      inputs,
      coords=None,
      mask=None,
      numerical_stabilizer=1e-9,
      train: bool = False,
  ):
      """Applies self attention on the input data.
  
      Args:
        inputs: Input tensor of shape [batch_size, num_points, feature_dim]
        coords: Input tensor of point positions of the shape [batch_size,
          num_points, 3]
        mask: Binary array of shape broadcastable to `inputs` tensor, indicating
          the positions for which self attention should be computed.
        numerical_stabilizer: takes into account stability when inputs are zero.
        train: Whether it is training or not.
  
      Returns:
        output: Tensor of shape [batch_size, num_points, feature_dim]
      """
      input_q = nn.Conv(
          self.out_channels,
          kernel_size=(self.kernel_size, self.kernel_size),
          use_bias=True,
      )(inputs)
      input_k = nn.Conv(
          self.out_channels,
          kernel_size=(self.kernel_size, self.kernel_size),
          use_bias=True,
      )(inputs)
      input_v = nn.Conv(
          self.out_channels,
          kernel_size=(self.kernel_size, self.kernel_size),
          use_bias=True,
      )(inputs)
  
      if (
          self.attention_fn_configs is None
          or self.attention_fn_configs["attention_kind"] == "regular"
      ):
          attention = jnp.einsum("...MC,...NC->...MN", input_q, input_k)
          if mask is not None:
              mask = nn.make_attention_mask(mask, mask)
              mask = mask.squeeze(axis=-3)
              big_neg = jnp.finfo(attention.dtype).min
              attention = jnp.where(mask, attention, big_neg)
          attention = nn.softmax(attention, axis=-1)
          attention = attention / (
              numerical_stabilizer
              + jnp.sum(attention, axis=1, keepdims=True)
          )
          output = jnp.einsum("...MN,...NC->...NC", attention, input_v)
      else:
          query = jnp.expand_dims(input_q, axis=-2)
          key = jnp.expand_dims(input_k, axis=-2)
          value = jnp.expand_dims(input_v, axis=-2)
          # TODO(kchoro): Include point cloud masking in performer attention
          if (
              self.attention_fn_configs["performer"]["masking_type"]
              == "nomask"
          ):
              output = (
                  performer.regular_performer_dot_product_attention(
                      query,
                      key,
                      value,
                      kernel_config=self.attention_fn_configs[
                          "performer"
                      ],
                  )
              )
          elif (
              self.attention_fn_configs["performer"]["masking_type"]
              == "fftmasked"
          ):
              toeplitz_params = self.param(
                  "toeplitz_params",
                  zeros,
                  (query.shape[-2], 2 * query.shape[-3] - 1),
              )
              output = performer.masked_performer_dot_product_attention(
                  query,
                  key,
                  value,
                  toeplitz_params=toeplitz_params,
                  kernel_config=self.attention_fn_configs["performer"],
              )
          elif (
              self.attention_fn_configs["performer"]["masking_type"]
              == "sharpmasked"
          ):
              toeplitz_params = self.param(
                  "toeplitz_params",
                  zeros,
                  (query.shape[-2], 5 * NUM_FT_PARAMS_PER_HEAD),
              )
              output = performer.sharp_masked_performer_dot_product_attention(
                  query,
                  key,
                  value,
                  coords,
                  toeplitz_params=toeplitz_params,
                  kernel_config=self.attention_fn_configs["performer"],
              )
          output = jnp.squeeze(output, axis=-2)
  
      output = (
          (inputs - output)
          if self.attention_type == "offset"
          else output
      )
      output = nn.Conv(
          self.out_channels,
          kernel_size=(self.kernel_size, self.kernel_size),
          use_bias=True,
      )(output)
      output = nn.BatchNorm(use_running_average=not train)(output)
      output = nn.relu(output)
      return output + inputs
  
imports: from typing import Any
  import flax.linen
  from flax.linen.initializers import zeros
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.projects.performer import performer
label_desc: Use image node_id 1 for calling the SelfAttentionLayer obj's underlying member method code with example usage: obj.__call__(inputs, coords, mask, numerical_stabilizer, train) and returns: unknown