def METHOD_NAME(device_type_mesh_map):
    """Configs corresponding mesh given test context.
    If runs on a CPU mesh, set virtual device on CPU.
    If runs on a GPU mesh, sets virtual device on GPU with proper memory
    limits.
    if runs on a TPU mesh, initializes TPU system.
    Args:
      device_type_mesh_map: A dictionary containing device_type -> mesh
        mapping.
    Returns:
      A properly configured mesh for use in test.
    """
    reset_context()
    def get_mesh(device_type):
        mesh = device_type_mesh_map.get(device_type, None)
        if mesh is None:
            dt = device_type
            raise ValueError(f"Requires a {dt} mesh to run test on {dt}.")
        return mesh
    mesh = None
    if tf.config.list_physical_devices("GPU"):
        mesh = get_mesh("GPU")
        reset_logical_devices("GPU", np.prod(mesh.shape()))
    else:
        mesh = get_mesh("CPU")
        reset_logical_devices("CPU", np.prod(mesh.shape()))
    context.ensure_initialized()
    return mesh