@@ -31,6 +31,7 @@
 #include <cmath>
 #include <cstdint>
 #include <utility>
+#include <vector>
 
 namespace o2::aod
 {
@@ -169,99 +170,116 @@ JETTAGGING_TABLES_DEF(Full, "F");
 // Tables for full inference of GNN-based b-jet tagger (including track origin, vertex predictions)
 namespace jetgnnindex
 {
-DECLARE_SOA_INDEX_COLUMN(JTrack, track);              //! index of the original track this in-jet track index corresponds to
-DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint8_t);  //! in-jet index i (0..N-1) of this constituent track, matching the order of jet.template tracks_as<AnyTracks>()
-DECLARE_SOA_COLUMN(EdgeIndex, edgeIndex, uint8_t);    //! unique in-jet edge index for the (i,j) constituent-track pair, i<j, triangular order i*(2N-i-1)/2+(j-i-1)
-DECLARE_SOA_DYNAMIC_COLUMN(TrackIds, trackIds, //! inverse of the triangular edge indexing above: for n constituent tracks, recovers the (i,j) in-jet track index pair (i<j) this edge connects
-                           [](uint8_t edgeIndex, int n) -> std::pair<uint8_t, uint8_t> {
-                             int i = 0;
-                             int offset = 0;
-                             while (offset + (n - 1 - i) <= static_cast<int>(edgeIndex)) {
-                               offset += (n - 1 - i);
-                               i++;
-                             }
-                             int j = i + 1 + (static_cast<int>(edgeIndex) - offset);
-                             return std::pair<uint8_t, uint8_t>{static_cast<uint8_t>(i), static_cast<uint8_t>(j)};
-                           });
+DECLARE_SOA_INDEX_COLUMN(JTrack, track);              //! index in JTracks of this GNN input track (not the parent AOD Tracks index)
+DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint32_t); //! node index in the GNN input after track selection/truncation; need not match the constituent-array position
+// DECLARE_SOA_COLUMN(EdgeIndex, edgeIndex, uint8_t);    //! unique in-jet edge index for the (i,j) constituent-track pair, i<j, triangular order i*(2N-i-1)/2+(j-i-1)
+// DECLARE_SOA_DYNAMIC_COLUMN(TrackIds, trackIds, //! inverse of the triangular edge indexing above: for n constituent tracks, recovers the (i,j) in-jet track index pair (i<j) this edge connects
+//                           [](uint8_t edgeIndex, int n) -> std::pair<uint8_t, uint8_t> {
+//                             int i = 0;
+//                             int offset = 0;
+//                             while (offset + (n - 1 - i) <= static_cast<int>(edgeIndex)) {
+//                               offset += (n - 1 - i);
+//                               i++;
+//                             }
+//                             int j = i + 1 + (static_cast<int>(edgeIndex) - offset);
+//                             return std::pair<uint8_t, uint8_t>{static_cast<uint8_t>(i), static_cast<uint8_t>(j)};
+//                           });
 } // namespace jetgnnindex
 
 namespace jetgnnlabel
 {
 DECLARE_SOA_COLUMN(Origin, origin, uint8_t);           //! trkOrigin from vertexClustering() (0: non-physical primary, 1: charm, 2: beauty, 3: primary vertex, 4: other secondary)
 DECLARE_SOA_COLUMN(VertexIndex, vertexIndex, uint8_t); //! trkVtxIndex from vertexClustering() (0: primary vertex, >0: secondary vertex, ranked by distance to PV)
-DECLARE_SOA_COLUMN(SameVertex, sameVertex, bool);      //! true if the (i,j) constituent-track pair of this edge share the same in-jet vertex cluster (same trkVtxIndex from vertexClustering())
+// DECLARE_SOA_COLUMN(SameVertex, sameVertex, bool);      //! true if the (i,j) constituent-track pair of this edge share the same in-jet vertex cluster (same trkVtxIndex from vertexClustering())
 } // namespace jetgnnlabel
 
 namespace jetgnnscore
 {
-DECLARE_SOA_COLUMN(ScoreML, scoreML, float); //! GNN prediction score (track/node level, or edge level symmetrised e.g. averaged from the N*N GNN edge matrix)
+DECLARE_SOA_COLUMN(ScoreML, scoreML, float);                        //! GNN prediction score (track/node level, or edge level symmetrised e.g. averaged from the N*N GNN edge matrix)
+DECLARE_SOA_COLUMN(OriginScores, originScores, std::vector<float>); //! one score per origin class, in jetgnnlabel::Origin order (0..4); producer defines logits/probabilities
 } // namespace jetgnnscore
 
-#define DECLARE_JET_GNNTRACKINDEX_TABLE(_jet_type_, _description_)         \
-  namespace _jet_type_##gnntrackindextable                                \
-  {                                                                       \
-    DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                           \
-  }                                                                      \
-  DECLARE_SOA_TABLE(_jet_type_##GNNTrack, "AOD", _description_,         \
-                    o2::soa::Index<>,                                   \
-                    _jet_type_##gnntrackindextable::_jet_type_##Id,     \
-                    jetgnnindex::JTrackId,                              \
-                    jetgnnindex::TrackIndex);                           \
-  namespace _jet_type_##gnntrackindices                                 \
-  {                                                                     \
-    DECLARE_SOA_ARRAY_INDEX_COLUMN(_jet_type_##GNNTrack, tracks);       \
-  }                                                                     \
-  DECLARE_SOA_TABLE(_jet_type_##GNNTrackIndices, "AOD", _description_ "I", \
+// GNNTrack has one row per evaluated (jet, track) pair. Its jet and track indices
+// identify the pair independently of the ordering in JetConstituents.
+// GNNTrackIndices is a positional jet extension: write one row for every jet in
+// Jets order, including an empty array for jets without evaluated tracks. Store
+// GNNTrack row indices (e.g. the backbone producer's lastIndex()), not JTrackIds.
+#define DECLARE_JET_GNNTRACKINDEX_TABLE(_jet_type_, _description_)                                                   \
+  namespace _jet_type_##gnntrackindextable                                                                           \
+  {                                                                                                                  \
+    DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                                                                       \
+  }                                                                                                                  \
+  DECLARE_SOA_TABLE(_jet_type_##GNNTrack, "AOD", _description_,                                                      \
+                    o2::soa::Index<>,                                                                                \
+                    _jet_type_##gnntrackindextable::_jet_type_##Id,                                                  \
+                    jetgnnindex::JTrackId,                                                                           \
+                    jetgnnindex::TrackIndex);                                                                        \
+  namespace _jet_type_##gnntrackindices                                                                              \
+  {                                                                                                                  \
+    DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_jet_type_##GNNTrack, tracksGnn, int32_t, _jet_type_##GNNTrack, "");         \
+  }                                                                                                                  \
+  DECLARE_SOA_TABLE(_jet_type_##GNNTrackIndices, "AOD", _description_ "I",                                           \
                     _jet_type_##gnntrackindices::_jet_type_##GNNTrackIds);
 
-#define DECLARE_JET_GNNEDGEINDEX_TABLE(_jet_type_, _description_)          \
-  namespace _jet_type_##gnnedgeindextable                                 \
-  {                                                                       \
-    DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                           \
-  }                                                                      \
-  DECLARE_SOA_TABLE(_jet_type_##GNNEdge, "AOD", _description_,          \
-                    o2::soa::Index<>,                                   \
-                    _jet_type_##gnnedgeindextable::_jet_type_##Id,      \
-                    jetgnnindex::EdgeIndex,                             \
-                    jetgnnindex::TrackIds<jetgnnindex::EdgeIndex>);     \
-  namespace _jet_type_##gnnedgeindices                                  \
-  {                                                                     \
-    DECLARE_SOA_ARRAY_INDEX_COLUMN(_jet_type_##GNNEdge, edges);         \
-  }                                                                     \
-  DECLARE_SOA_TABLE(_jet_type_##GNNEdgeIndices, "AOD", _description_ "I", \
-                    _jet_type_##gnnedgeindices::_jet_type_##GNNEdgeIds);
+// #define DECLARE_JET_GNNEDGEINDEX_TABLE(_jet_type_, _description_)          \
+//   namespace _jet_type_##gnnedgeindextable                                 \
+//   {                                                                       \
+//     DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                           \
+//   }                                                                      \
+//   DECLARE_SOA_TABLE(_jet_type_##GNNEdge, "AOD", _description_,          \
+//                     o2::soa::Index<>,                                   \
+//                     _jet_type_##gnnedgeindextable::_jet_type_##Id,      \
+//                     jetgnnindex::EdgeIndex,                             \
+//                     jetgnnindex::TrackIds<jetgnnindex::EdgeIndex>);     \
+//   namespace _jet_type_##gnnedgeindices                                  \
+//   {                                                                     \
+//     DECLARE_SOA_ARRAY_INDEX_COLUMN(_jet_type_##GNNEdge, edges);         \
+//   }                                                                     \
+//   DECLARE_SOA_TABLE(_jet_type_##GNNEdgeIndices, "AOD", _description_ "I", \
+//                     _jet_type_##gnnedgeindices::_jet_type_##GNNEdgeIds);
 
 // In-jet track/edge index backbone tables (own the in-jet ordering + the back-reference to the original track)
 DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedJet, "CJGNT");
 DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNT");
-DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedJet, "CJGNE");
-DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNE");
+// DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedJet, "CJGNE");
+// DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNE");
 
 #undef DECLARE_JET_GNNTRACKINDEX_TABLE
-#undef DECLARE_JET_GNNEDGEINDEX_TABLE
+// #undef DECLARE_JET_GNNEDGEINDEX_TABLE
 
 // Truth labels from vertexClustering() (MC detector level only). Thin tables: join positionally with
 // ChargedMCDetectorLevelJetGNNTrack / ChargedMCDetectorLevelJetGNNEdge (same producer loop, same row order).
 DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackOriginDef, "AOD", "CDJTOD", jetgnnlabel::Origin);
 DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetVertexDef, "AOD", "CDJTVD", jetgnnlabel::VertexIndex);
-DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex);
+// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex);
 
-// GNN ScoreML outputs, track (node) level, data and MC detector level. Thin tables: join positionally with
-// ChargedJetGNNTrack / ChargedMCDetectorLevelJetGNNTrack.
-DECLARE_SOA_TABLE(ChargedJetTrackTags, "AOD", "CJTRKT", jetgnnscore::ScoreML);
-DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackTags, "AOD", "CDJTRKT", jetgnnscore::ScoreML);
+// GNN track-origin outputs, data and MC detector level. Thin tables: write one
+// score vector per GNNTrack row, in the same row order as the backbone table.
+// The vector contains classes for this track, not scores for multiple tracks.
+DECLARE_SOA_TABLE(ChargedJetTrackTags, "AOD", "CJTRKT", jetgnnscore::OriginScores);
+DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackTags, "AOD", "CDJTRKT", jetgnnscore::OriginScores);
 
 // GNN ScoreML outputs, edge level, data and MC detector level. Thin tables: join positionally with
 // ChargedJetGNNEdge / ChargedMCDetectorLevelJetGNNEdge.
-DECLARE_SOA_TABLE(ChargedJetEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML);
-DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML);
-
-// Convenience joined types combining the backbone (in-jet ordering / original-track reference) with the thin
-// content tables, so `.template tracks_as<ChargedMCDetectorLevelJetTrackLabels>()` etc. can be used directly.
-using ChargedJetGNNTrackLabels = o2::soa::Join<ChargedJetGNNTrack, ChargedJetTrackTags>;
-using ChargedJetGNNEdgeLabels = o2::soa::Join<ChargedJetGNNEdge, ChargedJetEdgeTags>;
+// DECLARE_SOA_TABLE(ChargedJetEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML);
+// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML);
+
+// Prediction joins do not require MC truth tables. Subscribe to the prediction
+// join and JTracks in the consuming task to bind both index levels:
+//   for (auto const& prediction : jet.tracksGnn_as<ChargedJetGNNTrackPredictions>()) {
+//     auto track = prediction.track_as<JTracks>();
+//     auto scores = prediction.originScores();
+//   }
+using ChargedJetGNNTrackPredictions = o2::soa::Join<ChargedJetGNNTrack, ChargedJetTrackTags>;
+using ChargedMCDetectorLevelJetGNNTrackPredictions = o2::soa::Join<ChargedMCDetectorLevelJetGNNTrack, ChargedMCDetectorLevelJetTrackTags>;
+using ChargedJetsWithGNNTrackOrigins = o2::soa::Join<ChargedJets, ChargedJetConstituents, ChargedJetGNNTrackIndices>;
+using ChargedMCDetectorLevelJetsWithGNNTrackOrigins = o2::soa::Join<ChargedMCDetectorLevelJets, ChargedMCDetectorLevelJetConstituents, ChargedMCDetectorLevelJetGNNTrackIndices>;
+
+// Convenience joins including truth labels when available (MC detector level).
+using ChargedJetGNNTrackLabels = ChargedJetGNNTrackPredictions;
+// using ChargedJetGNNEdgeLabels = o2::soa::Join<ChargedJetGNNEdge, ChargedJetEdgeTags>;
 using ChargedMCDetectorLevelJetGNNTrackLabels = o2::soa::Join<ChargedMCDetectorLevelJetGNNTrack, ChargedMCDetectorLevelJetTrackOriginDef, ChargedMCDetectorLevelJetVertexDef, ChargedMCDetectorLevelJetTrackTags>;
-using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join<ChargedMCDetectorLevelJetGNNEdge, ChargedMCDetectorLevelJetEdgeDef, ChargedMCDetectorLevelJetEdgeTags>;
+// using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join<ChargedMCDetectorLevelJetGNNEdge, ChargedMCDetectorLevelJetEdgeDef, ChargedMCDetectorLevelJetEdgeTags>;
 
 } // namespace o2::aod