class_docstr: Processes the output of both `SparseArch` (sparse_features) and `DenseArch`
  (dense_features). Return Y*Z and the dense layer itself (all concatenated)
  where Y is the output of interaction branch 1 and Z is the output of interaction
  branch 2. Y and Z are of size Bx(F1xD) and Bx(DxF2) respectively for some F1 and F2.
  
  .. note::
  
      The dimensionality of the `dense_features` (D) is expected to match the
      dimensionality of the `sparse_features` so that the dot products between them
      can be computed.
      The output dimension of the 2 interaction branches should be a multiple
      of D.
  
  
  Args:
      num_sparse_features (int): F.
      interaction_branch1 (nn.Module): MLP module for the first branch of
          interaction layer
      interaction_branch2 (nn.Module): MLP module for the second branch of
          interaction layer
  
  Example::
  
      D = 3
      B = 10
      keys = ["f1", "f2"]
      F = len(keys)
      # Assume last layer of
      I1 = DenseArch(
          in_features= 3 * D + D,
          layer_sizes=[4*D, 4*D], # F1 = 4
          device=dense_device,
      )
      I2 = DenseArch(
          in_features= 3 * D + D,
          layer_sizes=[4*D, 4*D], # F2 = 4
          device=dense_device,
      )
      inter_arch = InteractionProjectionArch(
                      num_sparse_features=len(keys),
                      interaction_branch1 = I1,
                      interaction_branch2 = I2,
                  )
  
      dense_features = torch.rand((B, D))
      sparse_features = torch.rand((B, F, D))
  
      #  B X (D + F1 * F2)
      concat_dense = inter_arch(dense_features, sparse_features)
class_docstr_tok: ["Processes","the","output","of","both","`","SparseArch","`","(","sparse_features",")","and","`","DenseArch","`","(","dense_features",")",".","Return","Y","*","Z","and","the","dense","layer","itself","(","all","concatenated",")","where","Y","is","the","output","of","interaction","branch","1","and","Z","is","the","output","of","interaction","branch","2",".","Y","and","Z","are","of","size","Bx","(","F1xD",")","and","Bx","(","DxF2",")","respectively","for","some","F1","and","F2",".","..","note",":",":","The","dimensionality","of","the","`","dense_features","`","(","D",")","is","expected","to","match","the","dimensionality","of","the","`","sparse_features","`","so","that","the","dot","products","between","them","can","be","computed",".","The","output","dimension","of","the","2","interaction","branches","should","be","a","multiple","of","D.","Args",":","num_sparse_features","(","int",")",":","F.","interaction_branch1","(","nn.Module",")",":","MLP","module","for","the","first","branch","of","interaction","layer","interaction_branch2","(","nn.Module",")",":","MLP","module","for","the","second","branch","of","interaction","layer","Example",":",":","D","=","3","B","=","10","keys","=","[","``","f1","''",",","``","f2","''","]","F","=","len","(","keys",")","#","Assume","last","layer","of","I1","=","DenseArch","(","in_features=","3","*","D","+","D",",","layer_sizes=","[","4","*","D",",","4","*","D","]",",","#","F1","=","4","device=dense_device",",",")","I2","=","DenseArch","(","in_features=","3","*","D","+","D",",","layer_sizes=","[","4","*","D",",","4","*","D","]",",","#","F2","=","4","device=dense_device",",",")","inter_arch","=","InteractionProjectionArch","(","num_sparse_features=len","(","keys",")",",","interaction_branch1","=","I1",",","interaction_branch2","=","I2",",",")","dense_features","=","torch.rand","(","(","B",",","D",")",")","sparse_features","=","torch.rand","(","(","B",",","F",",","D",")",")","#","B","X","(","D","+","F1","*","F2",")","concat_dense","=","inter_arch","(","dense_features",",","sparse_features",")"]
docstr: Args:
      dense_features (torch.Tensor): an input tensor of size B X D.
      sparse_features (torch.Tensor): an input tensor of size B X F X D.
  
  Returns:
      torch.Tensor: an output tensor of size B X (D + F1 * F2)) where
      F1*D and F2*D are the output dimensions of the 2 interaction MLPs.
docstr_tok: ["Args",":","dense_features","(","torch.Tensor",")",":","an","input","tensor","of","size","B","X","D.","sparse_features","(","torch.Tensor",")",":","an","input","tensor","of","size","B","X","F","X","D.","Returns",":","torch.Tensor",":","an","output","tensor","of","size","B","X","(","D","+","F1","*","F2",")",")","where","F1","*","D","and","F2","*","D","are","the","output","dimensions","of","the","2","interaction","MLPs","."]
code: def forward(
      self, dense_features: torch.Tensor, sparse_features: torch.Tensor
  ) -> torch.Tensor:
      """
      Args:
          dense_features (torch.Tensor): an input tensor of size B X D.
          sparse_features (torch.Tensor): an input tensor of size B X F X D.
  
      Returns:
          torch.Tensor: an output tensor of size B X (D + F1 * F2)) where
          F1*D and F2*D are the output dimensions of the 2 interaction MLPs.
      """
      if self.F <= 0:
          return dense_features
      (B, D) = dense_features.shape
  
      combined_values = torch.cat(
          (dense_features.unsqueeze(1), sparse_features), dim=1
      )
  
      interaction_branch1_out = self.interaction_branch1(
          torch.reshape(combined_values, (B, -1))
      )
  
      interaction_branch2_out = self.interaction_branch2(
          torch.reshape(combined_values, (B, -1))
      )
  
      interactions = torch.bmm(
          interaction_branch1_out.reshape([B, -1, D]),
          interaction_branch2_out.reshape([B, D, -1]),
      )
      interactions_flat = torch.reshape(interactions, (B, -1))
  
      return torch.cat((dense_features, interactions_flat), dim=1)
  
imports: from typing import Dict, List, Optional, Tuple
  import torch
  from torch import nn
  from torchrec.datasets.utils import Batch
  from torchrec.modules.crossnet import LowRankCrossNet
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.mlp import MLP
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
label_desc: Use image node_id 2 for calling the InteractionProjectionArch obj's underlying member method code with example usage: obj.forward(dense_features, sparse_features) and returns: torch, dense_features