code: def test_softmax_with_bwd(self):
      device_mesh = self.build_device_mesh()
  
      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 params in test_list:
          softmax_dim, shard_dim = params
          x = torch.rand(
              8, 12, 16, device=self.device_type, requires_grad=True
          )
          self.assertTrue(x.requires_grad)
          local_y = torch.nn.functional.softmax(
              x, dim=softmax_dim, dtype=torch.float32
          ).sum()
          local_y.backward()
  
          dist_x = distribute_tensor(x, device_mesh, [Shard(shard_dim)])
          self.assertTrue(dist_x.requires_grad)
          if dims[softmax_dim] == dims[shard_dim]:
              with self.assertRaisesRegex(
                  Exception, "Cannot run .* on sharding dimension!$"
              ):
                  dist_softmax = dist_x.softmax(dim=softmax_dim)
          else:
              dist_softmax = dist_x.softmax(dim=softmax_dim)
              shard_dim = (
                  shard_dim + dist_x.ndim
                  if shard_dim < 0
                  else shard_dim
              )
              self.assertTrue(
                  dist_softmax.placements[0].is_shard(dim=shard_dim)
              )
              dist_y = dist_softmax.sum()
              dist_y = dist_y.redistribute(device_mesh, [Replicate()])
              self.assertEqual(dist_y.to_local(), local_y)
              self.assertIsNone(dist_x.grad)
              dist_y.backward()
              self.assertIsNotNone(dist_x.grad)
              self.assertEqual(dist_x.grad.full_tensor(), x.grad)
  
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 5 for calling the DistMathOpsTest obj's underlying member method code with example usage: obj.test_softmax_with_bwd() without return types