class_docstr: SelfAttention Layer
  
  Given $X\in \mathbb{R}^{n   imes in_feature}$, the attention is calculated by: $a=softmax(W_2tanh(W_1X))$, where
  $W_1 \in \mathbb{R}^{hidden         imes in_feature}$, $W_2 \in \mathbb{R}^{out_feature     imes hidden}$.
  The final output is $y=aX$ where $y \in \mathbb{R}^{n       imes out_feature}$.
  
  Parameters
  ----------
  in_features: int
      Dimension of input features
  out_features: int
      Dimension of output features
  hidden_size: int
      Dimension of hidden layer
class_docstr_tok: ["SelfAttention","Layer","Given","$","X\\in","\\mathbb","{","R","}","^","{","n","imes","in_feature","}","$",",","the","attention","is","calculated","by",":","$","a=softmax","(","W_2tanh","(","W_1X",")",")","$",",","where","$","W_1","\\in","\\mathbb","{","R","}","^","{","hidden","imes","in_feature","}","$",",","$","W_2","\\in","\\mathbb","{","R","}","^","{","out_feature","imes","hidden","}","$",".","The","final","output","is","$","y=aX","$","where","$","y","\\in","\\mathbb","{","R","}","^","{","n","imes","out_feature","}","$",".","Parameters","--","--","--","--","--","in_features",":","int","Dimension","of","input","features","out_features",":","int","Dimension","of","output","features","hidden_size",":","int","Dimension","of","hidden","layer"]
docstr: The forward function.
  
  Parameters
  ----------
  X: torch.Tensor
      input feature of shape $\mathbb{R}^{n       imes in_feature}$.
  
  Returns
  -------
  embedding: torch.Tensor
      The final embedding of shape $\mathbb{R}^{out_features      imes in_feature}$
  attention-matrix: torch.Tensor
      The attention matrix
docstr_tok: ["The","forward","function",".","Parameters","--","--","--","--","--","X",":","torch.Tensor","input","feature","of","shape","$","\\mathbb","{","R","}","^","{","n","imes","in_feature","}","$",".","Returns","--","--","--","-","embedding",":","torch.Tensor","The","final","embedding","of","shape","$","\\mathbb","{","R","}","^","{","out_features","imes","in_feature","}","$","attention-matrix",":","torch.Tensor","The","attention","matrix"]
code: def forward(self, X):
      """The forward function.
  
      Parameters
      ----------
      X: torch.Tensor
          input feature of shape $\mathbb{R}^{n \times in_feature}$.
  
      Returns
      -------
      embedding: torch.Tensor
          The final embedding of shape $\mathbb{R}^{out_features \times in_feature}$
      attention-matrix: torch.Tensor
          The attention matrix
      """
      x = torch.tanh(torch.matmul(self.w1, X.transpose(1, 0)))
      x = torch.matmul(self.w2, x)
      attn = torch.nn.functional.softmax(x, dim=-1)
      x = torch.matmul(attn, X)
      return x, attn
  
file_path: deepchem/deepchem/models/torch_models/attention.py
imports: import math
  from typing import Optional
label_desc: Use image node_id 3 for calling the SelfAttention obj's underlying member method code with example usage: obj.forward(X) and returns: x,  attn