code: def test_populate_weights(self) -> None:
      pw = PositionWeightedModule(max_feature_length=10)
  
      #     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]),
      )
      features = features.to_dict()
  
      jt = features["f1"]
      weighted_features = pw(jt)
  
      self.assertEqual(weighted_features.weights().size(), (3,))
  
      pw_f1_ref = torch.gather(
          pw.state_dict()["position_weight"], 0, torch.tensor([0, 1, 0])
      )
  
      pw_f1 = weighted_features.weights().detach()
      self.assertTrue(torch.allclose(pw_f1_ref, pw_f1))
  
      position_weighted_module_gm = symbolic_trace(pw)
      position_weighted_module_gm_script = torch.jit.script(
          position_weighted_module_gm
      )
  
      weighted_features_gm_script = position_weighted_module_gm_script(
          jt
      )
      torch.testing.assert_close(
          weighted_features.values(),
          weighted_features_gm_script.values(),
      )
      torch.testing.assert_close(
          weighted_features.lengths(),
          weighted_features_gm_script.lengths(),
      )
  
file_path: torchrec/torchrec/modules/tests/test_feature_processor_.py
imports: import unittest
  import torch
  from torchrec.distributed.utils import init_parameters
  from torchrec.fx.tracer import symbolic_trace
  from torchrec.modules.feature_processor_ import PositionWeightedModule, PositionWeightedModuleCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 1 for calling the PositionWeightedModuleTest obj's underlying member method code with example usage: obj.test_populate_weights() without return types