code: def test_populate_weights(self) -> None:
      position_weighted_module_collection = (
          PositionWeightedModuleCollection({"f1": 10, "f2": 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]),
      )
  
      fp_kjt = position_weighted_module_collection(features)
  
      position_weighted_module_collection_gm = symbolic_trace(
          position_weighted_module_collection
      )
      position_weighted_module_collection_gm_script = torch.jit.script(
          position_weighted_module_collection_gm
      )
      fp_kjt_gm_script = position_weighted_module_collection_gm_script(
          features
      )
  
      torch.testing.assert_close(
          fp_kjt.values(), fp_kjt_gm_script.values()
      )
      torch.testing.assert_close(
          fp_kjt.lengths(), fp_kjt_gm_script.lengths()
      )
      torch.testing.assert_close(
          fp_kjt.length_per_key(), fp_kjt_gm_script.length_per_key()
      )
  
      empty_kjt = KeyedJaggedTensor.from_lengths_sync(
          keys=[],
          values=torch.tensor([], dtype=torch.int32),
          lengths=torch.tensor([], dtype=torch.int32),
      )
  
      empty_fp_kjt = position_weighted_module_collection(empty_kjt)
      empty_fp_kjt_gm_script = (
          position_weighted_module_collection_gm_script(empty_kjt)
      )
  
      torch.testing.assert_close(
          empty_fp_kjt.values(), empty_fp_kjt_gm_script.values()
      )
      torch.testing.assert_close(
          empty_fp_kjt.lengths(), empty_fp_kjt_gm_script.lengths()
      )
      torch.testing.assert_close(
          empty_fp_kjt.length_per_key(),
          empty_fp_kjt_gm_script.length_per_key(),
      )
  
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 PositionWeightedCollectionModuleTest obj's underlying member method code with example usage: obj.test_populate_weights() without return types