class_docstr: Layout of a tensor.
  
  Attributes:
      device_mesh: the device mesh to store the tensor distributed.
      sharding_spec: the sharding specification to describe how the tensor is sharded.
      global_shape: the entire shape of the global tensor.
class_docstr_tok: ["Layout","of","a","tensor",".","Attributes",":","device_mesh",":","the","device","mesh","to","store","the","tensor","distributed",".","sharding_spec",":","the","sharding","specification","to","describe","how","the","tensor","is","sharded",".","global_shape",":","the","entire","shape","of","the","global","tensor","."]
code: def get_sharded_shape_per_device(self):
      sharded_shape = list(self.global_shape)
      for (
          dim,
          shard_list,
      ) in self.sharding_spec.dim_partition_dict.items():
          mesh_list = [
              self.device_mesh.shape[mesh_dim]
              for mesh_dim in shard_list
          ]
          shard_partitions = reduce(operator.mul, mesh_list, 1)
          assert (
              sharded_shape[dim] % shard_partitions == 0
          ), f"Cannot shard dimension {dim} into {shard_partitions} partitions."
          sharded_shape[dim] //= shard_partitions
      return torch.Size(sharded_shape)
  
file_path: colossalai/colossalai/tensor/d_tensor/layout.py
imports: import operator
  from functools import reduce
  import torch
  from colossalai.device.device_mesh import DeviceMesh
  from .misc import DuplicatedShardingDimensionError, ShardingNotDivisibleError
  from .sharding_spec import ShardingSpec
label_desc: Use image node_id 3 for calling the Layout obj's underlying member method code with example usage: obj.get_sharded_shape_per_device() and returns: torch