code: def test_baddbmm(self):
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      tensor = torch.rand(
          4, 4, 8, device=self.device_type, requires_grad=True
      )
      batch_1 = torch.rand(
          4, 4, 8, device=self.device_type, requires_grad=True
      )
      batch_2 = torch.rand(
          4, 8, 8, device=self.device_type, requires_grad=True
      )
  
      def test_placement_comb(
          tensor_placements: List[Placement],
          batch_1_placements: List[Placement],
          batch_2_placements: List[Placement],
          beta: int,
          alpha: int,
          batch_1_grad: Optional[torch.Tensor],
      ) -> None:
          tensor_dt = distribute_tensor(
              tensor, device_mesh, tensor_placements
          )
          batch_1_dt = distribute_tensor(
              batch_1, device_mesh, batch_1_placements
          )
          batch_2_dt = distribute_tensor(
              batch_2, device_mesh, batch_2_placements
          )
          dist_res = cast(
              DTensor,
              torch.baddbmm(
                  tensor_dt,
                  batch_1_dt,
                  batch_2_dt,
                  beta=beta,
                  alpha=alpha,
              ),
          ).redistribute(device_mesh, [Replicate()])
          dist_local_res = dist_res.to_local()
          assert not torch.isnan(local_result).any()
          assert not torch.isnan(dist_local_res).any()
          self.assertEqual(
              dist_local_res.detach(), local_result.detach()
          )
  
          # TODO: add test backward
          # grad_dist_res = torch.ones_like(dist_res)
          # dist_res.backward(grad_dist_res)
          # self.assertIsNotNone(batch_1_dt.grad)
          # batch_1_grad_local = batch_1_dt.grad.redistribute(
          #     device_mesh, [Replicate()]
          # ).to_local()
          # self.assertEqual(batch_1_grad_local, batch_1_grad)
  
      shard0_spec = Shard(0)
      shard1_spec = Shard(1)
      shard2_spec = Shard(2)
      replica_spec = Replicate()
      shard_specs = [
          shard0_spec,
          shard1_spec,
          shard2_spec,
          replica_spec,
      ]
      shard_specs_comb = list(
          itertools.product(shard_specs, shard_specs, shard_specs)
      )
      # If beta is 0, input tensor will be ignored
      numeric_params_comb = [
          (0.0, 0.5),  # zero-beta
          (0.8, 0.5),  # non-zero-beta
      ]
  
      for beta, alpha in numeric_params_comb:
          local_result = torch.baddbmm(
              tensor, batch_1, batch_2, beta=beta, alpha=alpha
          )
          grad_local_res = torch.ones_like(local_result)
          local_result.backward(grad_local_res)
          # test all combos
          for spec in shard_specs_comb:
              test_placement_comb(
                  [spec[0]],
                  [spec[1]],
                  [spec[2]],
                  beta,
                  alpha,
                  batch_1.grad,
              )
  
file_path: pytorch/test/distributed/_tensor/test_matrix_ops.py
imports: import itertools
  from typing import cast, List, Optional
  import torch
  from torch.distributed._tensor import DeviceMesh, distribute_tensor
  from torch.distributed._tensor.api import DTensor
  from torch.distributed._tensor.placement_types import _Partial, Placement, 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 6 for calling the DistMatrixOpsTest obj's underlying member method code with example usage: obj.test_baddbmm() without return types