code: def test_distribute_module(self):
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      # fully shard all linear modules on dim 0
      module_to_shard = MyModel(
          5 * self.world_size, 20, device=self.device_type
      )
      shard_spec = [Shard(0)]
  
      def shard_fn(name, module, device_mesh):
          if isinstance(module, nn.Linear):
              for name, param in module.named_parameters():
                  dist_param = torch.nn.Parameter(
                      distribute_tensor(param, device_mesh, shard_spec)
                  )
                  module.register_parameter(name, dist_param)
  
      sharded_module = distribute_module(
          module_to_shard, device_mesh, shard_fn
      )
      for param in sharded_module.parameters():
          self.assertIsInstance(param, DTensor)
          self.assertEqual(param.placements, shard_spec)
  
      replica_spec = [Replicate()]
      # fully replicate all modules without passing in partition_fn
      module_to_replicate = MyModel(5, 20, device=self.device_type)
      replica_module = distribute_module(
          module_to_replicate, device_mesh
      )
      for param in replica_module.parameters():
          self.assertIsInstance(param, DTensor)
          self.assertEqual(param.placements, replica_spec)
  
      # fully replicate all modules by passing in partition_fn
      def replicate_fn(name, module, device_mesh):
          if isinstance(module, nn.Linear):
              for name, param in module.named_parameters():
                  dist_param = torch.nn.Parameter(
                      distribute_tensor(
                          param, device_mesh, replica_spec
                      )
                  )
                  module.register_parameter(name, dist_param)
  
      module_to_replicate = MyModel(5, 20, device=self.device_type)
      replica_module = distribute_module(
          module_to_replicate, device_mesh, replicate_fn
      )
      for param in replica_module.parameters():
          self.assertIsInstance(param, DTensor)
          self.assertEqual(param.placements, replica_spec)
  
      # only shard part of module, and rest of module should be replicate
      def shard_fn(name, module, device_mesh):
          if isinstance(module, nn.Linear) and (
              name == "seq.0" or name == "seq.8"
          ):
              for name, param in module.named_parameters():
                  dist_param = torch.nn.Parameter(
                      distribute_tensor(param, device_mesh, shard_spec)
                  )
                  module.register_parameter(name, dist_param)
  
      module_to_distribute = MyModel(
          5 * self.world_size, 20, device=self.device_type
      )
      dist_module = distribute_module(
          module_to_distribute, device_mesh, shard_fn
      )
      for name, param in dist_module.named_parameters():
          self.assertIsInstance(param, DTensor)
          if name.startswith(("seq.0", "seq.8")):
              self.assertEqual(param.placements, shard_spec)
          else:
              self.assertEqual(param.placements, replica_spec)
  
file_path: pytorch/test/distributed/_tensor/test_api.py
imports: import torch
  import torch.nn
  from torch.distributed._tensor import DeviceMesh, distribute_module, distribute_tensor, DTensor, 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 DTensorAPITest obj's underlying member method code with example usage: obj.test_distribute_module() without return types