class_docstr: PS tables correspond to an EmbeddingCollection or EmbeddingBagCollection.
class_docstr_tok: ["PS","tables","correspond","to","an","EmbeddingCollection","or","EmbeddingBagCollection","."]
docstr: Create PSCollection for `sharded_module`, whose module path is `path`
  
  Args:
      path: module path of the sharded module.
      sharded_module: the sharded module.
      params_plan: the sharding plan of `sharded_module`.
      url: configuration for PS, e.g. redis://127.0.0.1:6379/?prefix=model.
      ps_config: config of the PS, currently only support setting chunk size.
  
  Return:
      PSCollection of the sharded module.
docstr_tok: ["Create","PSCollection","for","`","sharded_module","`",",","whose","module","path","is","`","path","`","Args",":","path",":","module","path","of","the","sharded","module",".","sharded_module",":","the","sharded","module",".","params_plan",":","the","sharding","plan","of","`","sharded_module","`",".","url",":","configuration","for","PS",",","e.g",".","redis",":","\/\/127.0.0.1:6379\/","?","prefix=model",".","ps_config",":","config","of","the","PS",",","currently","only","support","setting","chunk","size",".","Return",":","PSCollection","of","the","sharded","module","."]
code: def fromModule(
      path, sharded_module, params_plan, url, ps_config=None
  ):
      """
      Create PSCollection for `sharded_module`, whose module path is `path`
  
      Args:
          path: module path of the sharded module.
          sharded_module: the sharded module.
          params_plan: the sharding plan of `sharded_module`.
          url: configuration for PS, e.g. redis://127.0.0.1:6379/?prefix=model.
          ps_config: config of the PS, currently only support setting chunk size.
  
      Return:
          PSCollection of the sharded module.
      """
  
      state_dict = sharded_module.state_dict()
      optimizer_state_dict = (
          sharded_module.fused_optimizer.state_dict()["state"]
      )
      tensor_infos = {}
      for key, tensor in state_dict.items():
          # Here we use the fact that state_dict will be shape of
          # `embeddings.xxx.weight` or `embeddingbags.xxx.weight`
          if (
              len(key.split(".")) <= 1
              or key.split(".")[1] not in params_plan
          ):
              continue
          table_name = key.split(".")[1]
          param_plan = params_plan.pop(table_name)
          tensors = [tensor]
          # This is really hardcoded right now...
          optimizer_state = optimizer_state_dict[key]
          if f"{table_name}.momentum1" in optimizer_state:
              tensors.append(optimizer_state[f"{table_name}.momentum1"])
          if f"{table_name}.momentum2" in optimizer_state:
              tensors.append(optimizer_state[f"{table_name}.momentum2"])
          tensor_infos[table_name] = (param_plan, tensors)
  
      assert (
          len(params_plan) == 0
      ), f"There are sharded param not found, leaving: {params_plan}."
  
      if isinstance(url, str):
          collection_schema = url
      else:
          collection_schema = lambda table_name: url(path, table_name)
  
      return PSCollection(
          path, tensor_infos, collection_schema, ps_config
      )
  
file_path: torchrec/contrib/dynamic_embedding/src/torchrec_dynamic_embedding/ps.py
imports: import os
  from typing import Callable, Dict, List, Tuple, Union
  import torch
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torchrec.distributed.types import ParameterSharding
  from .tensor_list import TensorList
label_desc: Use image node_id 4 for calling the PSCollection obj's underlying member method code with example usage: obj.fromModule(path, sharded_module, params_plan, url, ps_config) and returns: PSCollection