def build_pathnet_graph(p_inputs, p_labels, p_task_id, pathnet, training):
  """Builds a PathNet graph, returns input placeholders and output tensors.

  Args:
    p_inputs: (tf.placeholder) placeholder for the input image.
    p_labels: (tf.placeholder) placeholder for the target labels.
    p_task_id: (tf.placeholder) placeholder for the task id.
    pathnet: (pn.PathNet) PathNet model to use.
    training: (bool) whether the graph is being created for training.

  Returns:
    A pair of (`train_step_op`, `out_logits`), where `train_step_op` is
    the training op, and `out_logits` is the final network output
    (classification logits).
  """
  end_state = pathnet({
      'in_tensor': p_inputs,
      'labels': p_labels,
      'task_id': p_task_id,
      'training': training
  })

  # The train op is the same for all tasks. Note that the last layer in
  # PathNet models contains heads that compute the task losses, with one head
  # per task. The head is chosen depending on `task_id`, so the loss for
  # the current task is always under `task_loss` key in `end_state`.
  train_step_op = end_state['train_op']
  out_logits = end_state['in_tensor']

  return train_step_op, out_logits