class_docstr: PositionWeightedProcessor represents a processor to apply position weight to a KeyedJaggedTensor.
  
  It can handle both unsharded and sharded input and output corresponding output
  
  Args:
      max_feature_lengths (Dict[str, int]): Dict of feature_lengths, the key is the feature_name and value is length.
      device (Optional[torch.device]): default compute device.
  
  Example::
  
      keys=["Feature0", "Feature1", "Feature2"]
      values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 3, 4, 5, 6, 7])
      lengths=torch.tensor([2, 0, 1, 1, 1, 3, 2, 3, 0])
      features = KeyedJaggedTensor.from_lengths_sync(keys=keys, values=values, lengths=lengths)
      pw = FeatureProcessorCollection(
          feature_processor_modules={key: PositionWeightedFeatureProcessor(max_feature_length=100) for key in keys}
      )
      result = pw(features)
      # result is
      # KeyedJaggedTensor({
      #     "Feature0": {
      #         "values": [[0, 1], [], [2]],
      #         "weights": [[1.0, 1.0], [], [1.0]]
      #     },
      #     "Feature1": {
      #         "values": [[3], [4], [5, 6, 7]],
      #         "weights": [[1.0], [1.0], [1.0, 1.0, 1.0]]
      #     },
      #     "Feature2": {
      #         "values": [[3, 4], [5, 6, 7], []],
      #         "weights": [[1.0, 1.0], [1.0, 1.0, 1.0], []]
      #     }
      # })
docstr: In unsharded or non-pipelined model, the input features both contain fp_feature
  and non_fp_features, and the output will filter out non_fp features
  In sharded pipelining model, the input features can only contain either none
  or all feature_processed features, since the input feature comes from the
  input_dist() of ebc which will filter out the keys not in the ebc. And the
  input size is same as output size
  
  Args:
      features (KeyedJaggedTensor): input features
  
  Returns:
      KeyedJaggedTensor
docstr_tok: ["In","unsharded","or","non-pipelined","model",",","the","input","features","both","contain","fp_feature","and","non_fp_features",",","and","the","output","will","filter","out","non_fp","features","In","sharded","pipelining","model",",","the","input","features","can","only","contain","either","none","or","all","feature_processed","features",",","since","the","input","feature","comes","from","the","input_dist","(",")","of","ebc","which","will","filter","out","the","keys","not","in","the","ebc",".","And","the","input","size","is","same","as","output","size","Args",":","features","(","KeyedJaggedTensor",")",":","input","features","Returns",":","KeyedJaggedTensor"]
code: def forward(self, features: KeyedJaggedTensor) -> KeyedJaggedTensor:
      """
      In unsharded or non-pipelined model, the input features both contain fp_feature
      and non_fp_features, and the output will filter out non_fp features
      In sharded pipelining model, the input features can only contain either none
      or all feature_processed features, since the input feature comes from the
      input_dist() of ebc which will filter out the keys not in the ebc. And the
      input size is same as output size
  
      Args:
          features (KeyedJaggedTensor): input features
  
      Returns:
          KeyedJaggedTensor
      """
      if is_fx_tracing():
          features_dict = features.to_dict()
          weighted_features_names: List[str] = []
          weighted_features_values: List[torch.Tensor] = []
          weighted_features_lengths: List[torch.Tensor] = []
          weighted_features_weights: List[torch.Tensor] = []
          for key, position_weight in self.position_weights.items():
              seq = torch.ops.fbgemm.offsets_range(
                  features_dict[key].offsets().long(),
                  torch.numel(features_dict[key].values()),
              )
              weighted_features_names.append(key)
              weighted_features_values.append(
                  features_dict[key].values()
              )
              weighted_features_lengths.append(
                  features_dict[key].lengths()
              )
              weighted_features_weights.append(
                  torch.gather(position_weight, dim=0, index=seq)
              )
          return KeyedJaggedTensor.from_lengths_sync(
              keys=weighted_features_names,
              values=torch.cat(weighted_features_values),
              lengths=torch.cat(weighted_features_lengths),
              weights=torch.cat(weighted_features_weights),
          )
      else:
          feature_names = features.keys()
          lengths = features.lengths()
          offsets = features.offsets()
          values = features.values()
          length_per_key = features.length_per_key()
          weights = features.weights_or_none()
          batch_size = features.stride()
  
          has_fp_id_list_feature = False
          has_normal_id_list_feature = False
  
          if weights is None:
              cat_seq = torch.ops.fbgemm.offsets_range(
                  offsets.long(), torch.numel(values)
              )
          else:
              # for row-wise sharding
              cat_seq = weights.long()
          seqs = torch.split(cat_seq, features.length_per_key())
  
          for feature_name in feature_names:
              if feature_name in self.max_feature_lengths:
                  has_fp_id_list_feature = True
              else:
                  has_normal_id_list_feature = True
  
          # in sharded pipelining model, the input features can only contain either none
          # or all feature_processed features, since the input feature comes from the
          # input_dist() of ebc which will filter out the keys not in the ebc
          # for the input features both contain fp_feature and normal_features, it could be
          # unsharded or non-pipelined sharded models
          if has_fp_id_list_feature:
              # for sharded pipeling
              if not has_normal_id_list_feature:
                  processed_features_weights: List[torch.Tensor] = []
                  for feature_index, feature_name in enumerate(
                      feature_names
                  ):
                      processed_weight = torch.gather(
                          self.position_weights[feature_name],
                          dim=0,
                          index=seqs[feature_index],
                      )
                      processed_features_weights.append(
                          processed_weight
                      )
                  fp_features = KeyedJaggedTensor(
                      keys=feature_names,
                      values=values,
                      weights=torch.cat(processed_features_weights),
                      lengths=lengths,
                      offsets=offsets,
                      stride=batch_size,
                      length_per_key=length_per_key,
                      offset_per_key=features.offset_per_key(),
                      index_per_key=features._key_indices(),
                  )
              # for unsharded or sharded non-pipeling
              else:
                  feature_values = values.split(length_per_key)
                  feature_lengths = lengths.split(batch_size)
                  processed_features_names: List[str] = []
                  processed_features_lengths: List[torch.Tensor] = []
                  processed_features_values: List[torch.Tensor] = []
                  processed_features_weights: List[torch.Tensor] = []
                  for feature_index, feature_name in enumerate(
                      feature_names
                  ):
                      if feature_name in self.max_feature_lengths:
                          feature_value = feature_values[feature_index]
                          feature_length = feature_lengths[
                              feature_index
                          ]
                          processed_weight = torch.gather(
                              self.position_weights[feature_name],
                              dim=0,
                              index=seqs[feature_index],
                          )
                          processed_features_names.append(feature_name)
                          processed_features_lengths.append(
                              feature_length
                          )
                          processed_features_values.append(
                              feature_value
                          )
                          processed_features_weights.append(
                              processed_weight
                          )
                  fp_features = KeyedJaggedTensor.from_lengths_sync(
                      keys=processed_features_names,
                      values=torch.cat(processed_features_values),
                      lengths=torch.cat(processed_features_lengths),
                      weights=torch.cat(processed_features_weights),
                  )
              return fp_features
          # normal id_list feature
          else:
              return features
  
imports: import abc
  from collections import OrderedDict
  from typing import Any, Dict, Iterator, List, Optional, Tuple
  import torch
  import torch.nn
  from torchrec.fx.tracer import is_fx_tracing
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the PositionWeightedProcessor obj's underlying member method code with example usage: obj.forward(features) and returns: KeyedJaggedTensor, fp_features, features