def filter_by_device_name(items, device_names, target_device_name):
  """Filter a list of items by device name.

  Args:
    items: A list of items to be filtered according to their corresponding
      device names.
    device_names: A list of the device names. Must have the same legnth
      as `items`.
    target_device_name: A `str` representing the desired device name.

  Returns:
    Filtered items from `items`.
  """
  assert len(items) == len(device_names)
  assert all(device_names), "device_names are not all non-empty strings"
  # Note: we use `endswith` instead of `==` for device-name filtering because
  # in some cases, the device names from kernel/op execution can have slightly
  # different values than the device names from
  # `distribution.extended.worker_devices`.
  return [items[i] for i, device_name in enumerate(device_names)
          if device_name.endswith(target_device_name)]