code: def test_rowwise_parallel_style(self):
      mesh = init_device_mesh(self.device_type, (self.world_size,))
  
      comm_mode = CommDebugMode()
      tensor = torch.rand(
          8,
          16 // self.world_size,
          device=self.device_type,
          requires_grad=True,
      )
      model = nn.Linear(16, 16, device=self.device_type)
  
      default_row_parallel = RowwiseParallel()
      with comm_mode:
          rowwise_mod = parallelize_module(
              deepcopy(model), mesh, default_row_parallel
          )
          out = rowwise_mod(tensor)
          # ensure output replicated
          self.assertEqual(out.shape, (8, 16))
          # allreduce in fwd
          self.assertEqual(
              comm_mode.get_comm_counts()[c10d_functional.all_reduce], 1
          )
          self.assertEqual(comm_mode.get_total_counts(), 1)
  
          out.sum().backward()
          # no op in bwd
          self.assertEqual(comm_mode.get_total_counts(), 1)
  
      sharded_row_parallel = RowwiseParallel(output_layouts=Shard(0))
      with comm_mode:
          rowwise_mod = parallelize_module(
              deepcopy(model), mesh, sharded_row_parallel
          )
          out = rowwise_mod(tensor)
          # ensure output replicated
          self.assertEqual(out.shape, (8 // self.world_size, 16))
          # reduce_scatter in fwd
          self.assertEqual(
              comm_mode.get_comm_counts()[
                  c10d_functional.reduce_scatter_tensor
              ],
              1,
          )
          self.assertEqual(comm_mode.get_total_counts(), 1)
  
          out.sum().backward()
          # allgather in bwd
          self.assertEqual(
              comm_mode.get_comm_counts()[
                  c10d_functional.all_gather_into_tensor
              ],
              1,
          )
          self.assertEqual(comm_mode.get_total_counts(), 2)
  
file_path: pytorch/test/distributed/tensor/parallel/test_tp_style.py
imports: from copy import deepcopy
  import torch
  import torch.nn
  from torch.distributed._tensor import Replicate, Shard, init_device_mesh
  from torch.distributed._tensor.debug import CommDebugMode
  from torch.distributed.tensor.parallel import parallelize_module
  from torch.distributed.tensor.parallel.style import ColwiseParallel, PrepareModuleInput, PrepareModuleOutput, RowwiseParallel
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms, NUM_DEVICES
label_desc: Use image node_id 3 for calling the TensorParallelStyleTest obj's underlying member method code with example usage: obj.test_rowwise_parallel_style() without return types