code: def test_weight_tying(self):
      class TestModule(torch.nn.Module):
          def __init__(self):
              super().__init__()
              # Initialize different weights for embedding and fc.
              torch.manual_seed(1)
              self.embedding = torch.nn.Embedding(16, 8)
              torch.manual_seed(2)
              self.fc = torch.nn.Linear(8, 16)
  
          def forward(self, x):
              return self.fc(self.embedding(x))
  
      model = TestModule().to(self.device_type)
      parallelize_plan = {
          "embedding": ColwiseParallel(),
          "fc": RowwiseParallel(),
      }
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      parallelize_module(model, device_mesh, parallelize_plan)
  
      input_size = [5]
      torch.manual_seed(0)
      inp = torch.randint(16, input_size, device=self.device_type)
  
      # Without weight tying.
      self.assertNotEqual(
          model.embedding.weight.to_local(), model.fc.weight.to_local()
      )
      output = model(inp)
      output.sum().backward()
      self.assertNotEqual(
          model.embedding.weight.grad.to_local(),
          model.fc.weight.grad.to_local(),
      )
      model.zero_grad()
  
      # With weight tying.
      model.fc.weight = model.embedding.weight
  
      self.assertEqual(model.embedding.weight, model.fc.weight)
      self.assertEqual(id(model.embedding.weight), id(model.fc.weight))
      output = model(inp)
      output.sum().backward()
      self.assertEqual(
          model.embedding.weight.grad, model.fc.weight.grad
      )
      self.assertEqual(
          id(model.embedding.weight.grad), id(model.fc.weight.grad)
      )
  
file_path: pytorch/test/distributed/tensor/parallel/test_tp_examples.py
imports: from copy import deepcopy
  import torch
  import torch.distributed
  from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard
  import torch.distributed._functional_collectives
  from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper, CheckpointImpl
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, PrepareModuleInput, RowwiseParallel
  from torch.distributed.tensor.parallel.input_reshard import input_reshard
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, MLPModule, ModelArgs, NUM_DEVICES, skip_unless_torch_gpu, Transformer, with_comms
label_desc: Use image node_id 7 for calling the DistTensorParallelExampleTest obj's underlying member method code with example usage: obj.test_weight_tying() and returns: self