code: def test_addmm_auto_redistribute(self):
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      shard0_spec = [Shard(0)]
      shard1_spec = [Shard(1)]
      replica_spec = [Replicate()]
  
      tensor_to_shard1 = torch.randn(12, 8, requires_grad=True)
      mat1 = distribute_tensor(
          tensor_to_shard1, device_mesh, shard1_spec
      )
      tensor_to_shard0 = torch.randn(8, 4, requires_grad=True)
      mat2 = distribute_tensor(
          tensor_to_shard0, device_mesh, shard0_spec
      )
      input_tensor = torch.randn(4, requires_grad=True)
      input = distribute_tensor(input_tensor, device_mesh, replica_spec)
  
      local_res = torch.addmm(
          input_tensor, tensor_to_shard1, tensor_to_shard0
      )
      dist_res = torch.addmm(input, mat1, mat2)
  
      # test if addmm output is a partial
      self.assertIsInstance(dist_res, DTensor)
      self.assertIsInstance(dist_res.placements[0], _Partial)
  
      # test if result is the same as tensor
      dist_local_res = dist_res.full_tensor()
      self.assertEqual(local_res, dist_local_res)
  
      # backward checks
      dist_local_res.sum().backward()
      local_res.sum().backward()
      self.assertIsNotNone(mat2.grad)
      self.assertEqual(mat2.grad.full_tensor(), tensor_to_shard0.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 2 for calling the DistMatrixOpsTest obj's underlying member method code with example usage: obj.test_addmm_auto_redistribute() without return types