The device of an operation.

    Note that only tf operations have device assignments.

    Args:
      operation_name: a string, name of an operation in the graph.

    Returns:
      a string or None, representing the device name.

def get_operation_device(self, operation_name):
    """The device of an operation.

    Note that only tf operations have device assignments.

    Args:
      operation_name: a string, name of an operation in the graph.

    Returns:
      a string or None, representing the device name.
    """
    operation = self._name_to_operation(operation_name)
    if isinstance(operation, tf.Operation):
      return operation.device
    else:  # mtf.Operation
      return None