code: def test_layer_norm_fwd(self):
      device_mesh = self.build_device_mesh()
  
      # NLP example from pytorch docs
      # https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
      batch, sentence_length, embedding_dim = 20, 5, 10
      x = torch.rand(
          batch, sentence_length, embedding_dim, device=self.device_type
      )
      norm_shape_idx_list = list(range(x.ndim))
      shard_dims = [-1, 0, 1, 2]
      elementwise_affine_list = [False, True]
      test_config_list = list(
          itertools.product(
              shard_dims, norm_shape_idx_list, elementwise_affine_list
          )
      )
  
      # normalized shape is a torch.Size object
      for shard_dim, norm_idx, elementwise_affine in test_config_list:
          normalized_shape = x.shape[norm_idx:]
          layer_norm = torch.nn.LayerNorm(
              normalized_shape,
              elementwise_affine=elementwise_affine,
              device=self.device_type,
          )
          layer_norm_local = copy.deepcopy(layer_norm).to(
              self.device_type
          )
  
          def _replicate_fn(name, module, device_mesh):
              for name, param in module.named_parameters():
                  if name in ["weight", "bias"]:
                      param_dist = torch.nn.Parameter(
                          distribute_tensor(
                              param, device_mesh, [Replicate()]
                          )
                      )
                      module.register_parameter(name, param_dist)
  
          layer_norm_dist = distribute_module(
              layer_norm, device_mesh, _replicate_fn
          )
  
          x_local = x
          x_dist = distribute_tensor(x, device_mesh, [Shard(shard_dim)])
  
          y_local = layer_norm_local(x_local)
          # make sure that forward layer norm does not introduce extra collectives
          comm_mode = CommDebugMode()
          with comm_mode:
              y_dist = layer_norm_dist(x_dist)
  
          self.assertEqual(
              comm_mode.get_total_counts(),
              0,
              f"comm count={comm_mode.get_total_counts()}, "
              f"shard_dim={shard_dim}, norm_shape={normalized_shape}, elem_affine={elementwise_affine}",
          )
  
          from torch.distributed._tensor.placement_types import (
              TensorMeta,
          )
  
          dtensor_meta = y_dist._spec.tensor_meta
          assert isinstance(dtensor_meta, TensorMeta)
          # make sure the right shape in sharding prop
          self.assertEqual(y_local.shape, dtensor_meta.shape)
          self.assertEqual(y_local, y_dist.full_tensor())
  
file_path: pytorch/test/distributed/_tensor/test_math_ops.py
imports: import copy
  import itertools
  import torch
  from torch.distributed._tensor import DeviceMesh, distribute_module, distribute_tensor
  from torch.distributed._tensor.debug import CommDebugMode
  from torch.distributed._tensor.ops.utils import is_tensor_partial
  from torch.distributed._tensor.placement_types import Replicate, Shard
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, skip_unless_torch_gpu, with_comms
label_desc: Use image node_id 7 for calling the DistMathOpsTest obj's underlying member method code with example usage: obj.test_layer_norm_fwd() without return types