code: def forward(self, features: KeyedJaggedTensor) -> KeyedJaggedTensor:
      scores_list = []
      for feature_name in features.keys():
          jt = features[feature_name]
          scores = jt.weights()
          scores = torch.clamp(
              scores,
              min=self.feature_min,
              max=self.feature_max,
          )
          indices = torch.floor(torch.pow(scores, self.feature_pow))
          indices = indices.to(torch.int32)
          scores = torch.index_select(self.w, 0, indices)
          scores_list.append(scores)
  
      return KeyedJaggedTensor(
          keys=features.keys(),
          values=features.values(),
          weights=torch.cat(scores_list)
          if scores_list
          else features.weights_or_none(),
          lengths=features.lengths(),
          stride=features.stride(),
      )
  
imports: import math
  import unittest
  from typing import Dict, List, Tuple
  import hypothesis.strategies
  import torch
  from hypothesis import given, settings
  from torchrec import EmbeddingBagConfig, EmbeddingCollection, EmbeddingConfig, KeyedJaggedTensor
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, ShardingType
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
  from torchrec.distributed.planner.shard_estimators import EmbeddingPerfEstimator, EmbeddingStorageEstimator
  from torchrec.distributed.quant_embeddingbag import QuantFeatureProcessedEmbeddingBagCollectionSharder
  from torchrec.distributed.quant_state import sharded_tbes_weights_spec, WeightSpec
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, placement, row_wise
  from torchrec.distributed.test_utils.infer_utils import assert_close, assert_weight_spec, create_cw_min_partition_constraints, create_test_model, KJTInputWrapper, model_input_to_forward_args, model_input_to_forward_args_kjt, prep_inputs, quantize, quantize_fpebc, shard_qebc, shard_qec, TestModelInfo, TestQuantEBCSharder
  from torchrec.distributed.test_utils.test_model import ModelInput
  from torchrec.distributed.types import ShardingEnv, ShardingPlan
  from torchrec.fx import symbolic_trace
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.feature_processor_ import FeatureProcessorsCollection
  from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
label_desc: Use image node_id 2 for calling the TimeGapPoolingCollectionModule obj's underlying member method code with example usage: obj.forward(features) and returns: KeyedJaggedTensor