code: def test_replicate_to_partial(self):
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      local_tensor = torch.randn(
          12, 3, device=self.device_type, requires_grad=True
      )
      partial_spec = _Partial()
      replica_spec = Replicate()
      # 1) test replicate -> partial forward
      replica_tensor = distribute_tensor(
          local_tensor, device_mesh, [replica_spec]
      )
      with self.assertRaisesRegex(
          RuntimeError, "Can not redistribute to _Partial"
      ):
          partial_tensor = replica_tensor.redistribute(
              device_mesh, [partial_spec]
          )
  
      from torch.distributed._tensor.redistribute import Redistribute
  
      partial_tensor = Redistribute.apply(
          replica_tensor, device_mesh, [partial_spec]
      )
      self.assertEqual(partial_tensor.size(), local_tensor.size())
      # test it successfully zero out the contents on other ranks
      self.assertEqual(
          replica_tensor.to_local() / self.world_size,
          partial_tensor.to_local(),
      )
  
      # replicate to partial on sub groups
      local_tensor = torch.randn(12, 3, device=self.device_type)
      device_mesh = DeviceMesh(
          self.device_type,
          torch.arange(self.world_size).reshape(
              self.world_size // 2, 2
          ),
      )
      # 1) test replicate -> partial on 2d-mesh subgroups
      replica_tensor = distribute_tensor(
          local_tensor, device_mesh, [replica_spec, replica_spec]
      )
      partial_tensor = Redistribute.apply(
          replica_tensor, device_mesh, [partial_spec, partial_spec]
      )
      self.assertEqual(partial_tensor.size(), local_tensor.size())
  
      self.assertEqual(
          replica_tensor.to_local() / self.world_size,
          partial_tensor.to_local(),
      )
  
file_path: pytorch/test/distributed/_tensor/test_redistribute.py
imports: import itertools
  import torch
  from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor
  from torch.distributed._tensor.placement_types import _Partial, Replicate, Shard
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
label_desc: Use image node_id 5 for calling the RedistributeTest obj's underlying member method code with example usage: obj.test_replicate_to_partial() without return types