code: def forward(self, features: KeyedJaggedTensor) -> KeyedJaggedTensor:
      if features.weights_or_none() is None:
          cat_seq = torch.ops.fbgemm.offsets_range(
              features.offsets().long(), torch.numel(features.values())
          )
      else:
          # for row-wise sharding
          cat_seq = features.weights().long()
      seqs = torch.split(cat_seq, features.length_per_key())
      weights_list = []
      for key, seq in zip(features.keys(), seqs):
          if key in self.max_feature_lengths:
              weights_list.append(
                  torch.gather(
                      self.position_weights[key], dim=0, index=seq
                  )
              )
          else:
              weights_list.append(
                  self._dummy_weights[: self.max_feature_lengths[key]]
              )
      weights = torch.cat(weights_list)
  
      return KeyedJaggedTensor(
          keys=features.keys(),
          values=features.values(),
          weights=weights,
          lengths=features.lengths(),
          offsets=features.offsets(),
          stride=features.stride(),
          length_per_key=features.length_per_key(),
      )
  
file_path: torchrec/torchrec/distributed/grouped_position_weighted.py
imports: from collections import OrderedDict
  from typing import Any, Dict, Iterator, Optional, Tuple
  import torch
  import torch.nn
  from torchrec.distributed.embedding_types import BaseGroupedFeatureProcessor
  from torchrec.distributed.utils import append_prefix
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the GroupedPositionWeightedModule obj's underlying member method code with example usage: obj.forward(features) and returns: KeyedJaggedTensor