code: def test_flat_data(self):
      state_dict = {
          "a": torch.rand(10),
          "b": [1, 2, 3],
      }
  
      a = find_state_dict_object(state_dict, MetadataIndex("a"))
      self.assertEqual(a, state_dict["a"])
      a = find_state_dict_object(state_dict, MetadataIndex("a", [0]))
      self.assertEqual(a, state_dict["a"])
      a = find_state_dict_object(
          state_dict, MetadataIndex("a", index=99)
      )
      self.assertEqual(a, state_dict["a"])
  
      b = find_state_dict_object(state_dict, MetadataIndex("b"))
      self.assertEqual(b, state_dict["b"])
      b = find_state_dict_object(
          state_dict, MetadataIndex("b", index=1)
      )
      self.assertEqual(b, state_dict["b"])
  
      with self.assertRaisesRegex(ValueError, "FQN"):
          find_state_dict_object(state_dict, MetadataIndex("c"))
      with self.assertRaisesRegex(ValueError, "ShardedTensor"):
          find_state_dict_object(state_dict, MetadataIndex("b", [1]))
  
file_path: pytorch/test/distributed/checkpoint/test_utils.py
imports: import sys
  import torch
  from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor, ShardedTensorMetadata, ShardMetadata
  from torch.distributed._shard.sharded_tensor.metadata import TensorProperties
  from torch.distributed.checkpoint.metadata import MetadataIndex
  from torch.distributed.checkpoint.utils import find_state_dict_object
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN, TestCase
  from torch.testing._internal.distributed.distributed_utils import with_fake_comms
label_desc: Use image node_id 4 for calling the TestMedatadaIndex obj's underlying member method code with example usage: obj.test_flat_data() without return types