class_docstr: This is the Factorization Machine module, mentioned in the `DeepFM paper
  <https://arxiv.org/pdf/1703.04247.pdf>`_:
  
  This module does not cover the end-end functionality of the published paper.
  Instead, it covers only the FM part of the publication, and is used to learn
  2nd-order feature interactions.
  
  To support modeling flexibility, we customize the key components as different from
  the public paper:
      We change the input from raw sparse features to embeddings of the features.
      This allows flexibility in embedding dimensions and the number of embeddings,
      as long as all embedding tensors have the same batch size.
  
  The general architecture of the module is like::
  
                              1 x 10                  output
                               /|\
                                |                     pass into `dense_module`
                                |
                              1 x 90
                               /|\
                                |                     concat
                                |
                      1 x 20, 1 x 30, 1 x 40          list of embeddings
  
  Example::
  
      batch_size = 3
      # the input embedding are in torch.Tensor of [batch_size, num_embeddings, embedding_dim]
      input_embeddings = [
          torch.randn(batch_size, 2, 64),
          torch.randn(batch_size, 2, 32),
      ]
      fm = FactorizationMachine()
      output = fm(embeddings=input_embeddings)
class_docstr_tok: ["This","is","the","Factorization","Machine","module",",","mentioned","in","the","`","DeepFM","paper","<","https",":","\/\/arxiv.org\/pdf\/1703.04247.pdf",">","`","_",":","This","module","does","not","cover","the","end-end","functionality","of","the","published","paper",".","Instead",",","it","covers","only","the","FM","part","of","the","publication",",","and","is","used","to","learn","2nd-order","feature","interactions",".","To","support","modeling","flexibility",",","we","customize","the","key","components","as","different","from","the","public","paper",":","We","change","the","input","from","raw","sparse","features","to","embeddings","of","the","features",".","This","allows","flexibility","in","embedding","dimensions","and","the","number","of","embeddings",",","as","long","as","all","embedding","tensors","have","the","same","batch","size",".","The","general","architecture","of","the","module","is","like",":",":","1","x","10","output","\/|\\","|","pass","into","`","dense_module","`","|","1","x","90","\/|\\","|","concat","|","1","x","20",",","1","x","30",",","1","x","40","list","of","embeddings","Example",":",":","batch_size","=","3","#","the","input","embedding","are","in","torch.Tensor","of","[","batch_size",",","num_embeddings",",","embedding_dim","]","input_embeddings","=","[","torch.randn","(","batch_size",",","2",",","64",")",",","torch.randn","(","batch_size",",","2",",","32",")",",","]","fm","=","FactorizationMachine","(",")","output","=","fm","(","embeddings=input_embeddings",")"]
docstr: Args:
      embeddings (List[torch.Tensor]):
          The list of all embeddings (e.g. dense, common_sparse,
          specialized_sparse, embedding_features, raw_embedding_features) in the
          shape of::
  
              (batch_size, num_embeddings, embedding_dim)
  
          For the ease of operation, embeddings that have the same embedding
          dimension have the option to be stacked into a single tensor. For
          example, when we have 1 trained embedding with dimension=32, 5 native
          embeddings with dimension=64, and 3 dense features with dimension=16, we
          can prepare the embeddings list to be the list of::
  
              tensor(B, 1, 32) (trained_embedding with num_embeddings=1, embedding_dim=32)
              tensor(B, 5, 64) (native_embedding with num_embeddings=5, embedding_dim=64)
              tensor(B, 3, 16) (dense_features with num_embeddings=3, embedding_dim=32)
  
          NOTE:
              `batch_size` of all input tensors need to be identical.
  
  Returns:
      torch.Tensor: output of fm with flattened and concatenated `embeddings` as input. Expected to be [B, 1].
docstr_tok: ["Args",":","embeddings","(","List","[","torch.Tensor","]",")",":","The","list","of","all","embeddings","(","e.g",".","dense",",","common_sparse",",","specialized_sparse",",","embedding_features",",","raw_embedding_features",")","in","the","shape","of",":",":","(","batch_size",",","num_embeddings",",","embedding_dim",")","For","the","ease","of","operation",",","embeddings","that","have","the","same","embedding","dimension","have","the","option","to","be","stacked","into","a","single","tensor",".","For","example",",","when","we","have","1","trained","embedding","with","dimension=32",",","5","native","embeddings","with","dimension=64",",","and","3","dense","features","with","dimension=16",",","we","can","prepare","the","embeddings","list","to","be","the","list","of",":",":","tensor","(","B",",","1",",","32",")","(","trained_embedding","with","num_embeddings=1",",","embedding_dim=32",")","tensor","(","B",",","5",",","64",")","(","native_embedding","with","num_embeddings=5",",","embedding_dim=64",")","tensor","(","B",",","3",",","16",")","(","dense_features","with","num_embeddings=3",",","embedding_dim=32",")","NOTE",":","`","batch_size","`","of","all","input","tensors","need","to","be","identical",".","Returns",":","torch.Tensor",":","output","of","fm","with","flattened","and","concatenated","`","embeddings","`","as","input",".","Expected","to","be","[","B",",","1","]","."]
code: def forward(
      self,
      embeddings: List[torch.Tensor],
  ) -> torch.Tensor:
      """
      Args:
          embeddings (List[torch.Tensor]):
              The list of all embeddings (e.g. dense, common_sparse,
              specialized_sparse, embedding_features, raw_embedding_features) in the
              shape of::
  
                  (batch_size, num_embeddings, embedding_dim)
  
              For the ease of operation, embeddings that have the same embedding
              dimension have the option to be stacked into a single tensor. For
              example, when we have 1 trained embedding with dimension=32, 5 native
              embeddings with dimension=64, and 3 dense features with dimension=16, we
              can prepare the embeddings list to be the list of::
  
                  tensor(B, 1, 32) (trained_embedding with num_embeddings=1, embedding_dim=32)
                  tensor(B, 5, 64) (native_embedding with num_embeddings=5, embedding_dim=64)
                  tensor(B, 3, 16) (dense_features with num_embeddings=3, embedding_dim=32)
  
              NOTE:
                  `batch_size` of all input tensors need to be identical.
  
      Returns:
          torch.Tensor: output of fm with flattened and concatenated `embeddings` as input. Expected to be [B, 1].
      """
  
      # flatten each embedding to be [B, N, D] -> [B, N*D], then cat them all on dim=1
      fm_input = _get_flatten_input(embeddings)
      sum_of_input = torch.sum(fm_input, dim=1, keepdim=True)
      sum_of_square = torch.sum(
          fm_input * fm_input, dim=1, keepdim=True
      )
      square_of_sum = sum_of_input * sum_of_input
      cross_term = square_of_sum - sum_of_square
      cross_term = (
          torch.sum(cross_term, dim=1, keepdim=True) * 0.5
      )  # [B, 1]
      return cross_term
  
file_path: torchrec/torchrec/modules/deepfm.py
imports: from typing import List
  import torch
  from torch import nn
  from torch.fx import wrap
label_desc: Use image node_id 2 for calling the FactorizationMachine obj's underlying member method code with example usage: obj.forward(embeddings) and returns: cross_term