# Interface defining how the Tensor should be placed on a set of devices. 
class PlacementSpec(object):
   pass
   
class DevicePlacement(PlacementSpec):
   # Device where a shard should be placed on in 
   # RemoteDevice format: https://github.com/pytorch/pytorch/issues/46554
   self.device: str
       
class SingleSplitShardingSpec(PlacementSpec):
    # List of placement of shards for the tensor should be sharded. Each 
    # entry in the List refers to one shard and defines where that shard 
    # should be placed.
    self.shard_placement: List[PlacementSpec]
    
    # A single integer indicating the dimension along which the tensor should be 
    # split. Could be dimension names as well following: 
    # https://pytorch.org/docs/stable/named_tensor.html
    self.split_dimension: int/str

    # Optional parameter to set the size of each shard (for uneven splits). 
    # If not specified, the split is even. Must be same length as self.devices.
    self.shard_sizes: List[int]
          
# Let's say we have a large tensor and we'd like to shard it row-wise across
# 4 GPUs on a single host. The sharding spec would be as follows:
spec = SingleSplitShardingSpec()
spec.devices = [
    "rank:0/cuda:0", "rank:1/cuda:1", "rank:2/cuda:2", "rank:3/cuda:3"
]
spec.sharding_def.split_dimension = 0

class MeshShardingSpec(PlacementSpec):
    # Multidimensional array with the same number of dimensions as the tensor 
    # to be sharded. Each element of the array represents a ShardPlacement 
    # where a particular partition of the tensor should be placed. The array 
    # represents how the tensor should be partitioned across multiple 
    # placement options. The number of elements in the array represents the 
    # total number of partitions.
    self.shard_placement: ShardPlacement.    
    
# Let's say we have a tensor of shape [4 x 8] as follows:
# 1 1 2 2 3 3 4 4
# 1 1 2 2 3 3 4 4
# 5 5 6 6 7 7 8 8
# 5 5 6 6 7 7 8 8
# And our, shard_placement array is [2 x 4] as follows:
[
    ["rank:0/cuda:0", "rank:0/cuda:1", "rank:0/cuda:2", "rank:0/cuda:3"],
    ["rank:1/cuda:4", "rank:1/cuda:5", "rank:1/cuda:6", "rank:1/cuda:7"],
] 

# Then each device gets a [2 x 2] slice of the tensor, where cuda:0 gets:
# 1 1
# 1 1
#, cuda:1 gets:
# 2 2
# 2 2
#, cuda:4 gets,
# 5 5
# 5 5
# and so on....

torch.distributed.sharding_spec.single_split(
    split_dimension, 
    shard_placement: List[PlacementSpec],
    shard_sizes = None,
) -> SingleSplitShardingSpec

# Example:
torch.distributed.sharding_spec.single_split(
    0, ["rank:0/cuda:0", "rank:1/cuda:1"])

torch.distributed.sharding_spec.mesh(
    shard_placement: Array[str]
) -> MeshShardingSpec

# Example:
torch.distributed.sharding_spec.mesh(
    [
        ["rank:0/cuda:0", "rank:0/cuda:1", "rank:0/cuda:2", "rank:0/cuda:3"],
        ["rank:1/cuda:4", "rank:1/cuda:5", "rank:1/cuda:6", "rank:1/cuda:7"],
    ]
)

CompositeShardingSpec(PlacementSpec):
    # Provides a mapping from "parameter name" to appropriate 
    # ShardingSpec/PlacementSpec
    param_placement: Dict[str, PlacementSpec]

# Helper API
# Accepts **kwargs of parameter name to PlacementSpec
torch.distributed.sharding_spec.composite(**kwargs)

# Example 1: Shard a module by placing parameters on different devices

class MyModel(nn.Module):
    def __init__(self):
        self.net1 = nn.Linear(20, 20)
        self.net2 = nn.Linear(20, 20)
        self.net3 = nn.Linear(20, 20)
        
torch.distributed.sharding_spec.composite(
    net1="cuda:0", net2="cuda:1", net3="cuda:0",
)

# Example 2: Apply sharding spec to submodules

from torch.distributed.sharding_spec import (composite, single_split)

class MyModel(nn.Module):
    def __init__(self):
        self.net1 = nn.Linear(20, 20)
        self.net2 = nn.Linear(20, 20)
        self.net3 = nn.Linear(20, 20)
        
# Spec for net1 and net2
composite_spec = composite(weight=single_split(0, ["cuda:0", "cuda:1"]))

# Spec for entire module
composite(
    net1=composite_spec, 
    net2=composite_spec, 
    net3=single_split(0, ["cuda:2", "cuda:3"])
)

# Similar to torch.empty, if needed we can have additional creation ops like 
# torch.ones/torch.zeros etc.
# This is done in SPMD fashion and needs to be called on all ranks. Each rank 
# will instantiate its local shard based on the ShardingSpec given.
torch.distributed.sharded_tensor.empty(
    *size, *, *sharding_spec**: ShardingSpec*, dtype=None, 
    requires_grad=False, pin_memory=False, names=None) → ShardedTensor

class MyModel(nn.Module):
    def __init__(self):
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 10)
        
model = MyModel()

# SPMD called on all ranks, rank 0 is used as the authoritative source of the 
# data and appropriate shards are broadcast to each rank.
# model.fc1.weight would be replaced with a ShardedTensor on each rank.
torch.distributed.shard_tensor(model.fc1.weight, sharding_spec)

# Need to be called in SPMD fashion on all ranks. Reshards the ShardedTensor 
# on which this method is called.
sharded_tensor = torch.distributed.sharded_tensor.empty(...)

sharded_tensor.reshard(
    sharding_spec: ShardingSpec
) -> ShardedTensor

model = my_sharded_model()

# On all ranks in SPMD fashion:

state_dict = my_sharded_model.state_dict()
# Save the state_dict on all ranks.
my_custom_save_model(state_dict, storage)

# Load the model from a checkpoint
state_dict = custom_load_state_dict(storage)
model = custom_init_model(state_dict)

# Now, if we'd like to modify sharding before training:
model.sharded_tensor1.reshard(sharding_spec)

# Using torch.save/torch.load:

model = my_sharded_model()

# On all ranks in SPMD fashion:

state_dict = my_sharded_model.state_dict()
# Save the state_dict on all ranks.
torch.save(state_dict, file)

# Load the model from a checkpoint
model.load_state_dict(torch.load(file))

# Returns a consolidated state_dict on rank 0 after collecting all shards of 
# all ShardedTensors.
torch.distributed.collect_state_dict(model.state_dict()) -> state_dict

# The state dict can be then loaded on a single rank, non-sharded weights 
# replicated and custom sharding spec can be applied to Tensors to shard them 
# using "torch.distributed.shard_tensor".

# Retrieves the Tensor representing the local shard of the ShardedTensor
sharded_tensor.local_shard() -> Tensor

# Retrieves the sharding metadata for the Tensor, which outlines how the 
# Tensor is sharded across all devices.
sharded_tensor.sharding_metadata() -> ShardedTensorMetadata

class ShardedTensorMetadata:
    self.shards: List[ShardMetadata]

# Metadata for a single shard.
class ShardMetadata:
    # Dimensions of the tensor representing this shard.
    self.dims: List[int],
    # Device that this shard is placed on (in RemoteDevice format:  https://github.com/pytorch/pytorch/issues/46554).
    self.device: str,
    # Starting offsets for each dimension of this shard in the 
    # ShardedTensor. Should be same size as self.dims. self.offsets 
    # combined with self.dims accurately identifies which partition of the 
    # Tensor this shard represents.
    self.offsets: List[int],

torch.distributed.shard_model(model: nn.Module, spec: PlacementSpec)