docstr: Decompose rewards of the last step to all individual steps.
  
  This model arch accepts sparse features and is similar to / inspired by
  the model in "Deep Learning Recommendation Model for Personalization and
  Recommendation Systems" (https://arxiv.org/abs/1906.00091)
  
  The model arch can be described as below:
  
  
                       last_layer_activation
                                 ^
                            overall arch
                                 ^
        -----------interaction arch (2D + 2F + F choose 2) -------
        ^                        ^                               ^
  state_dense_out(D)      action_dense_out(D)              sparse_out(F*D)
        ^                        ^                               ^
  state_dense_arch         action_dense_arch                sparse arch
        ^                        ^                               ^
    state_dense             action_dense             state_sparse / action_sparse
  
  
  , where:
  D: last layer of dense_sizes (equal to sparse features' embedding_dim)
  F: number of total sparse features (from both state and action and from both
      id-list and id-score-list features)
  Interaction arch returns a concatenation of
      (1) and the dense layers itself,
      (2) the dot product of each sparse embedding with the output of the dense arch,
      (3) the pairwise dot product of each sparse embedding pair,
docstr_tok: ["Decompose","rewards","of","the","last","step","to","all","individual","steps",".","This","model","arch","accepts","sparse","features","and","is","similar","to","\/","inspired","by","the","model","in","``","Deep","Learning","Recommendation","Model","for","Personalization","and","Recommendation","Systems","''","(","https",":","\/\/arxiv.org\/abs\/1906.00091",")","The","model","arch","can","be","described","as","below",":","last_layer_activation","^","overall","arch","^","--","--","--","--","--","-interaction","arch","(","2D","+","2F","+","F","choose","2",")","--","--","--","-","^","^","^","state_dense_out","(","D",")","action_dense_out","(","D",")","sparse_out","(","F","*","D",")","^","^","^","state_dense_arch","action_dense_arch","sparse","arch","^","^","^","state_dense","action_dense","state_sparse","\/","action_sparse",",","where",":","D",":","last","layer","of","dense_sizes","(","equal","to","sparse","features","'","embedding_dim",")","F",":","number","of","total","sparse","features","(","from","both","state","and","action","and","from","both","id-list","and","id-score-list","features",")","Interaction","arch","returns","a","concatenation","of","(","1",")","and","the","dense","layers","itself",",","(","2",")","the","dot","product","of","each","sparse","embedding","with","the","output","of","the","dense","arch",",","(","3",")","the","pairwise","dot","product","of","each","sparse","embedding","pair",","]
code: def __init__(
      self,
      state_dense_dim: int,
      action_dense_dim: int,
      dense_sizes: List[int],
      dense_activations: List[str],
      overall_sizes: List[int],
      overall_activations: List[str],
      embedding_bag_collection: EmbeddingBagCollection,
      use_batch_norm: bool = False,
      use_layer_norm: bool = False,
  ):
      """
      Decompose rewards of the last step to all individual steps.
  
      This model arch accepts sparse features and is similar to / inspired by
      the model in "Deep Learning Recommendation Model for Personalization and
      Recommendation Systems" (https://arxiv.org/abs/1906.00091)
  
      The model arch can be described as below:
  
  
                           last_layer_activation
                                     ^
                                overall arch
                                     ^
            -----------interaction arch (2D + 2F + F choose 2) -------
            ^                        ^                               ^
      state_dense_out(D)      action_dense_out(D)              sparse_out(F*D)
            ^                        ^                               ^
      state_dense_arch         action_dense_arch                sparse arch
            ^                        ^                               ^
        state_dense             action_dense             state_sparse / action_sparse
  
  
      , where:
      D: last layer of dense_sizes (equal to sparse features' embedding_dim)
      F: number of total sparse features (from both state and action and from both
          id-list and id-score-list features)
      Interaction arch returns a concatenation of
          (1) and the dense layers itself,
          (2) the dot product of each sparse embedding with the output of the dense arch,
          (3) the pairwise dot product of each sparse embedding pair,
  
      """
      super().__init__()
      self.validate_parameters(
          dense_sizes,
          dense_activations,
          overall_sizes,
          overall_activations,
          embedding_bag_collection,
      )
  
      self.state_dense_arch = create_dense_arch(
          state_dense_dim,
          dense_sizes,
          dense_activations,
          use_batch_norm,
          use_layer_norm,
      )
      self.action_dense_arch = create_dense_arch(
          action_dense_dim,
          dense_sizes,
          dense_activations,
          use_batch_norm,
          use_layer_norm,
      )
      # sparse arch will be shared for state sparse features and action sparse features
      self.sparse_arch = SparseArch(embedding_bag_collection)
  
      # Overall arch
      F = sum(
          [
              len(conf.feature_names)
              for conf in embedding_bag_collection.embedding_bag_configs()
          ]
      )
      D = dense_sizes[-1]
      self.F = F
      self.D = D
      sparse_feature_names = []
      for conf in embedding_bag_collection.embedding_bag_configs():
          sparse_feature_names.extend(conf.feature_names)
  
      try:
          self.inter_arch_sparse_and_state_dense = InteractionArch(
              F,
          )
          self.inter_arch_sparse_and_action_dense = InteractionArch(
              F,
          )
      except TypeError:
          # HACK: in torchrec OSS version (0.1.0), InteractionArch
          # only accepts a list of sparse feature names as the input
          # pyre-ignore
          self.inter_arch_sparse_and_state_dense = InteractionArch(
              sparse_feature_names=sparse_feature_names
          )
          # pyre-ignore
          self.inter_arch_sparse_and_action_dense = InteractionArch(
              sparse_feature_names=sparse_feature_names
          )
  
      interaction_output_dim = 2 * D + 2 * F + F * (F - 1) // 2
      self.overall_arch = create_dense_arch(
          interaction_output_dim,
          overall_sizes,
          overall_activations,
          use_batch_norm,
          use_layer_norm,
      )
  
imports: import logging
  from typing import List
  import torch
  import torch.nn
  from reagent.core import types
  from reagent.core.torch_utils import split_sequence_keyed_jagged_tensor
  from reagent.models.base import ModelBase
  from reagent.models.fully_connected_network import ACTIVATION_MAP
  from reagent.models.synthetic_reward import _gen_mask
  from torchrec import EmbeddingBagCollection
  from torchrec.models.dlrm import InteractionArch, SparseArch
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 1 to create a new SingleStepSyntheticSparseArchRewardNet object from inherited base classes: nn with example: obj = SingleStepSyntheticSparseArchRewardNet(state_dense_dim, action_dense_dim, dense_sizes, dense_activations, overall_sizes, overall_activations, embedding_bag_collection, use_batch_norm, use_layer_norm)