Add jet features after GNN

Useful for

  • Xbb for adding e.g. substructure info
  • Adding softmuon variables which @dkobylia may be interested in for SET/SMT