code: def test_populate_weights_PositionWeightedProcessor(self) -> None:
      features_max_length = {"f1": 10, "f2": 3}
      pw = PositionWeightedProcessor(features_max_length)
  
      #     0       1        2  <-- batch
      # 0   [0,1] None    [2]
      # 1   [3]    [4]    [5,6,7]
      # ^
      # feature
      features = KeyedJaggedTensor.from_offsets_sync(
          keys=["f1", "f2"],
          values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
          offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
      )
  
      weighted_features = pw(features)
  
      self.assertEqual(weighted_features["f1"].weights().size(), (3,))
      self.assertEqual(weighted_features["f2"].weights().size(), (5,))
  
      pw_f1_ref = torch.gather(
          pw.state_dict()["position_weights.f1"],
          0,
          torch.tensor([0, 1, 0]),
      )
      pw_f1 = weighted_features["f1"].weights().detach()
      self.assertTrue(torch.allclose(pw_f1_ref, pw_f1))
  
      pw_f2_ref = torch.gather(
          pw.state_dict()["position_weights.f2"],
          0,
          torch.tensor([0, 0, 0, 1, 2]),
      )
      pw_f2 = weighted_features["f2"].weights().detach()
      self.assertTrue(torch.allclose(pw_f2_ref, pw_f2))
  
file_path: torchrec/torchrec/modules/tests/test_feature_processor.py
imports: import unittest
  import torch
  from torchrec.fx import Tracer
  from torchrec.modules.feature_processor import PositionWeightedModule, PositionWeightedProcessor
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 3 for calling the PositionWeightedModuleTest obj's underlying member method code with example usage: obj.test_populate_weights_PositionWeightedProcessor() without return types