code: def test_softmax_fwd(self):
      device_mesh = self.build_device_mesh()
  
      x = torch.rand(8, 12, 16, device=self.device_type)
      dims = range(3)  # used to convert -1 to the actual dim
      softmax_dims = [-1, 0, 1, 2]
      shard_dims = [-1, 0, 1, 2]
      test_list = list(itertools.product(softmax_dims, shard_dims))
  
      for softmax_dim, shard_dim in test_list:
          local_y = torch.nn.functional.softmax(
              x, dim=softmax_dim, dtype=torch.float32
          )
          dist_x = distribute_tensor(x, device_mesh, [Shard(shard_dim)])
          if dims[shard_dim] == dims[softmax_dim]:
              with self.assertRaisesRegex(
                  Exception, "Cannot run .* on sharding dimension!$"
              ):
                  dist_y = torch.nn.functional.softmax(
                      dist_x, dim=softmax_dim, dtype=torch.float32
                  )
          else:
              dist_y = torch.nn.functional.softmax(
                  dist_x, dim=softmax_dim, dtype=torch.float32
              )
              shard_dim = (
                  shard_dim + dist_y.ndim
                  if shard_dim < 0
                  else shard_dim
              )
              self.assertTrue(
                  dist_y.placements[0].is_shard(dim=shard_dim)
              )
              self.assertEqual(dist_y.full_tensor(), local_y)
  
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 4 for calling the DistMathOpsTest obj's underlying member method code with example usage: obj.test_softmax_fwd() without return types