def METHOD_NAME(self):
    super().METHOD_NAME()
    global_ids = test_util.create_device_ids_array((2, 2))
    local_device_ids = np.ravel(global_ids).tolist()
    mesh_dict = {
        "CPU": dtensor.Mesh(
            ["X", "Y"],
            global_ids,
            local_device_ids,
            test_util.create_device_list((2, 2), "CPU"),
        )
    }
    self.mesh = self.configTestMesh(mesh_dict)