docstr: PS table of an embedding table.
  
  Args:
      table_name: name of the table.
      tensors: tensors of the table, the first one is the parameter tensor, others are
          tenors of optimizers, e.g. for Adam, it will be [weight, m, v].
      url: url of the PS.
docstr_tok: ["PS","table","of","an","embedding","table",".","Args",":","table_name",":","name","of","the","table",".","tensors",":","tensors","of","the","table",",","the","first","one","is","the","parameter","tensor",",","others","are","tenors","of","optimizers",",","e.g",".","for","Adam",",","it","will","be","[","weight",",","m",",","v","]",".","url",":","url","of","the","PS","."]
code: def __init__(
      self,
      table_name: str,
      tensors: Union[List[torch.Tensor], List[ShardedTensor]],
      url: str,
      chunk_size: str,
  ):
      """
      PS table of an embedding table.
  
      Args:
          table_name: name of the table.
          tensors: tensors of the table, the first one is the parameter tensor, others are
              tenors of optimizers, e.g. for Adam, it will be [weight, m, v].
          url: url of the PS.
      """
      shards = torch.classes.tde.LocalShardList()
      num_optimizer_stats = len(tensors)
      if isinstance(tensors[0], ShardedTensor):
          # Here we assume the shard metadata of optimizer state and weight are the same.
          for i, shard in enumerate(tensors[0].local_shards()):
              local_tensors = [
                  tensor.local_shards()[i].tensor for tensor in tensors
              ]
              shards.append(
                  shard.metadata.shard_offsets[0],
                  shard.metadata.shard_offsets[1],
                  shard.metadata.shard_sizes[0],
                  shard.metadata.shard_sizes[1],
                  TensorList(local_tensors).tensor_list,
              )
              # This assumes all shard have the same column size.
              col_size = shard.tensor.shape[1]
      elif isinstance(tensors[0], torch.Tensor):
          shards.append(
              0,
              0,
              tensors[0].shape[0],
              tensors[0].shape[1],
              TensorList(tensors).tensor_list,
          )
          col_size = tensors[0].shape[1]
      self._ps = torch.classes.tde.PS(
          table_name,
          shards,
          col_size,
          num_optimizer_stats,
          url,
          chunk_size,
      )
  
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 1 to create a new PS object  with example: obj = PS(table_name, tensors, url, chunk_size)