class_docstr: This is the `DeepFM module <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 deep component of the publication. It is used to learn
  high-order feature interactions. If low-order feature interactions should
  be learnt, please use `FactorizationMachine` module instead, which will share
  the same embedding input of this module.
  
  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. It allows flexibility in embedding dimensions and the
    number of embeddings, as long as all embedding tensors have the same batch size.
  
  * On top of the public paper, we allow users to customize the hidden layer to be any
    module, not limited to just MLP.
  
  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
  
  Args:
      dense_module (nn.Module):
          any customized module that can be used (such as MLP) in DeepFM. The
          `in_features` of this module must be equal to the element counts. For
          example, if the input embedding is `[randn(3, 2, 3), randn(3, 4, 5)]`, the
          `in_features` should be: 2*3+4*5.
  
  Example::
  
      import torch
      from torchrec.fb.modules.deepfm import DeepFM
      from torchrec.fb.modules.mlp import LazyMLP
      batch_size = 3
      output_dim = 30
      # the input embedding are a torch.Tensor of [batch_size, num_embeddings, embedding_dim]
      input_embeddings = [
          torch.randn(batch_size, 2, 64),
          torch.randn(batch_size, 2, 32),
      ]
      dense_module = nn.Linear(192, output_dim)
      deepfm = DeepFM(dense_module=dense_module)
      deep_fm_output = deepfm(embeddings=input_embeddings)
class_docstr_tok: ["This","is","the","`","DeepFM","module","<","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","deep","component","of","the","publication",".","It","is","used","to","learn","high-order","feature","interactions",".","If","low-order","feature","interactions","should","be","learnt",",","please","use","`","FactorizationMachine","`","module","instead",",","which","will","share","the","same","embedding","input","of","this","module",".","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",".","It","allows","flexibility","in","embedding","dimensions","and","the","number","of","embeddings",",","as","long","as","all","embedding","tensors","have","the","same","batch","size",".","*","On","top","of","the","public","paper",",","we","allow","users","to","customize","the","hidden","layer","to","be","any","module",",","not","limited","to","just","MLP",".","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","Args",":","dense_module","(","nn.Module",")",":","any","customized","module","that","can","be","used","(","such","as","MLP",")","in","DeepFM",".","The","`","in_features","`","of","this","module","must","be","equal","to","the","element","counts",".","For","example",",","if","the","input","embedding","is","`","[","randn","(","3",",","2",",","3",")",",","randn","(","3",",","4",",","5",")","]","`",",","the","`","in_features","`","should","be",":","2","*","3+4","*","5",".","Example",":",":","import","torch","from","torchrec.fb.modules.deepfm","import","DeepFM","from","torchrec.fb.modules.mlp","import","LazyMLP","batch_size","=","3","output_dim","=","30","#","the","input","embedding","are","a","torch.Tensor","of","[","batch_size",",","num_embeddings",",","embedding_dim","]","input_embeddings","=","[","torch.randn","(","batch_size",",","2",",","64",")",",","torch.randn","(","batch_size",",","2",",","32",")",",","]","dense_module","=","nn.Linear","(","192",",","output_dim",")","deepfm","=","DeepFM","(","dense_module=dense_module",")","deep_fm_output","=","deepfm","(","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 `dense_module` with flattened and concatenated `embeddings` as input.
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","`","dense_module","`","with","flattened","and","concatenated","`","embeddings","`","as","input","."]
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 `dense_module` with flattened and concatenated `embeddings` as input.
      """
  
      # flatten each embedding to be [B, N, D] -> [B, N*D], then cat them all on dim=1
      deepfm_input = _get_flatten_input(embeddings)
      deepfm_output = self.dense_module(deepfm_input)
      return deepfm_output
  
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 DeepFM obj's underlying member method code with example usage: obj.forward(embeddings) and returns: deepfm_output