code: def testModuleEviction(self):
      init_dist()
      rank = dist.get_rank()
      device = torch.device(f"cuda:{rank}")
      torch.cuda.set_device(device)
  
      embedding_dim = 8
      configs = [
          EmbeddingConfig(
              name="AB",
              num_embeddings=4,
              embedding_dim=embedding_dim,
              feature_names=["A", "B"],
              # to check reinit
              weight_init_min=2,
              weight_init_max=2,
          ),
      ]
      model = EmbeddingCollection(
          tables=configs, device=torch.device("meta")
      )
      model = DMP(module=model, device=device)
  
      plan = model.plan
      sharded_modules = _get_sharded_modules_recursive(
          model.module, "", plan
      )
      sharded_module, params_plan = sharded_modules[""]
      ps_collection = PSCollection.fromModule(
          "", sharded_module, params_plan, "memory://"
      )
  
      transformer_collection = IDTransformerCollection(
          configs, ps_collection=ps_collection
      )
  
      # extract weight manually
      weight = (
          model.module.state_dict()["embeddings.AB.weight"]
          .local_shards()[0]
          .tensor
      )
      self.assertTrue(
          weight.shape[0] == 4 and weight.shape[1] == embedding_dim
      )
  
      # init to zero
      weight[:] = torch.zeros_like(weight)
  
      global_kjt_1 = KeyedJaggedTensor(
          keys=["A", "B"],
          values=torch.tensor([1, 2, 3, 4, 1, 2]),
          lengths=torch.tensor([4, 2]),
      )
      cache_kjt, fetch_handles = transformer_collection.transform(
          global_kjt_1
      )
      for handle in fetch_handles:
          handle.wait()
      embedding = model(cache_kjt.to(device))
      self.assertTrue(
          torch.all(
              cache_kjt.values() == torch.tensor([0, 1, 2, 3, 0, 1])
          )
      )
      self.assertTrue(
          torch.allclose(
              embedding["A"].values(),
              torch.zeros_like(embedding["A"].values()),
          )
      )
      self.assertTrue(
          torch.allclose(
              embedding["B"].values(),
              torch.zeros_like(embedding["B"].values()),
          )
      )
  
      global_kjt_2 = KeyedJaggedTensor(
          keys=["A", "B"],
          values=torch.tensor([1, 2, 1, 5, 5, 6]),
          lengths=torch.tensor([3, 3]),
      )
  
      # this will evict 3 and 4
      cache_kjt, fetch_handles = transformer_collection.transform(
          global_kjt_2
      )
      for handle in fetch_handles:
          handle.wait()
      embedding = model(cache_kjt.to(device))
      self.assertTrue(
          torch.all(
              cache_kjt.values() == torch.tensor([0, 1, 0, 2, 2, 3])
          )
      )
      self.assertTrue(
          torch.allclose(
              embedding["A"].values(),
              torch.zeros_like(embedding["A"].values()),
          )
      )
      self.assertTrue(
          torch.allclose(
              embedding["B"].values(),
              2 * torch.ones_like(embedding["B"].values()),
          )
      )
  
      # update to one
      weight[:] = torch.ones_like(weight)
      global_kjt_3 = KeyedJaggedTensor(
          keys=["A", "B"],
          values=torch.tensor([1, 2, 1, 4, 3, 4]),
          lengths=torch.tensor([3, 3]),
      )
      cache_kjt, fetch_handles = transformer_collection.transform(
          global_kjt_3
      )
      for handle in fetch_handles:
          handle.wait()
      embedding = model(cache_kjt.to(device))
      self.assertTrue(
          torch.all(
              cache_kjt.values() == torch.tensor([0, 1, 0, 2, 3, 2])
          )
      )
      self.assertTrue(
          torch.allclose(
              embedding["A"].values(),
              torch.ones_like(embedding["A"].values()),
          )
      )
      self.assertTrue(
          torch.allclose(
              embedding["B"].values(),
              torch.zeros_like(embedding["B"].values()),
          )
      )
  
file_path: torchrec/contrib/dynamic_embedding/tests/test_ps_collection.py
imports: import unittest
  import torch
  import torch.distributed
  import torch.nn
  from torchrec import EmbeddingCollection, EmbeddingConfig, KeyedJaggedTensor
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec_dynamic_embedding.id_transformer_collection import IDTransformerCollection
  from torchrec_dynamic_embedding.ps import PSCollection
  from torchrec_dynamic_embedding.utils import _get_sharded_modules_recursive
  from utils import init_dist, register_memory_io
label_desc: Use image node_id 2 for calling the TestPSCollection obj's underlying member method code with example usage: obj.testModuleEviction() without return types