local_index_per_graph (tasks/utils.py, added in #87) regenerates per-graph indices from bincount(batch). Both predict_steps export these as the bus and idx identity columns in every parquet — predictions and now embeddings.
These are positions, not IDs. Correctness depends on datakit emitting buses in order and the wrapper's groupby preserving it. If that ever changes, the exported IDs silently mislabel rows.
Buses are at least validated (powergrid_hetero_dataset.py:162 asserts against the real bus column). Generators are not — gen_index is set from the positional DataFrame index (reset_index()), is never stored on the graph, and has no assert. Both spots carry TODOs saying to store real IDs in the tensors.
Fix: carry the true bus/gen IDs as tensor fields from the dataset onward and export those instead of regenerated positions; the ordering assert can then go away.
Refs #87
local_index_per_graph(tasks/utils.py, added in #87) regenerates per-graph indices frombincount(batch). Bothpredict_steps export these as thebusandidxidentity columns in every parquet — predictions and now embeddings.These are positions, not IDs. Correctness depends on datakit emitting buses in order and the wrapper's groupby preserving it. If that ever changes, the exported IDs silently mislabel rows.
Buses are at least validated (
powergrid_hetero_dataset.py:162asserts against the realbuscolumn). Generators are not —gen_indexis set from the positional DataFrame index (reset_index()), is never stored on the graph, and has no assert. Both spots carry TODOs saying to store real IDs in the tensors.Fix: carry the true bus/gen IDs as tensor fields from the dataset onward and export those instead of regenerated positions; the ordering assert can then go away.
Refs #87