def METHOD_NAME(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)
            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)
            dist_x_grad = dist_x.grad.redistribute(device_mesh, [Replicate()])
            self.assertEqual(dist_x_grad.to_local(), x.grad)