code: def test_bmm(self):
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      mat1 = torch.rand(
          4, 8, 4, device=self.device_type, requires_grad=True
      )
      mat2 = torch.rand(
          4, 4, 8, device=self.device_type, requires_grad=True
      )
      local_result = torch.bmm(mat1, mat2)
      grad_local_res = torch.ones_like(local_result)
      local_result.backward(grad_local_res)
  
      def test_placement_comb(
          placements1: List[Placement],
          placements2: List[Placement],
      ) -> None:
          mat1_dt = distribute_tensor(mat1, device_mesh, placements1)
          mat2_dt = distribute_tensor(mat2, device_mesh, placements2)
          dist_res = cast(
              DTensor, torch.bmm(mat1_dt, mat2_dt)
          ).redistribute(device_mesh, [Replicate()])
          dist_local_res = dist_res.to_local()
          self.assertEqual(dist_local_res, local_result)
  
          # test backward
          # TODO: figure out (replicate, shard1) fail on backward
          # it generates a different grad shape
          grad_dist_res = torch.ones_like(dist_res)
          dist_res.backward(grad_dist_res)
          self.assertIsNotNone(mat1_dt.grad)
          mat1_dt_grad = cast(DTensor, mat1_dt.grad)
          mat1_grad_local = mat1_dt_grad.redistribute(
              device_mesh, [Replicate()]
          ).to_local()
          self.assertEqual(mat1_grad_local, mat1.grad)
  
      shard0_spec = Shard(0)
      shard1_spec = Shard(1)
      shard2_spec = Shard(2)
      replica_spec = Replicate()
      placement_specs = [
          shard0_spec,
          shard1_spec,
          shard2_spec,
          replica_spec,
      ]
      shard_specs_comb = list(
          itertools.product(placement_specs, placement_specs)
      )
  
      # tests that currently pass
      for spec in shard_specs_comb:
          test_placement_comb([spec[0]], [spec[1]])
  
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 7 for calling the DistMatrixOpsTest obj's underlying member method code with example usage: obj.test_bmm() without return types