remove_nodes doesn't update graph_indicator on batched graphs #716

remove_nodes doesn't update graph_indicator, so on a batched graph the result keeps the old indicator (one entry per original node). The graph looks fine when printed, but unbatch / getgraph on it then throw a BoundsError.

MWE

using GNNGraphs, MLUtils

g = MLUtils.batch([rand_graph(3, 4, ndata = rand(Float32, 2, 3)),
                   rand_graph(4, 6, ndata = rand(Float32, 2, 4))])

g2 = remove_nodes(g, [1])

g2.num_nodes        # 6
g2.graph_indicator  # [1, 1, 1, 2, 2, 2, 2]  -> still 7 entries, should be [1, 1, 2, 2, 2, 2]

MLUtils.unbatch(g2) # BoundsError: attempt to access 2×6 Matrix{Float32} at index [1:2, 4:7]
getgraph(g2, 2)     # BoundsError: attempt to access 2×6 Matrix{Float32} at index [1:2, [4, 5, 6, 7]]

Node features, edges and the node count are all updated correctly, only graph_indicator is passed through unchanged in remove_nodes(g::GNNGraph{<:COO_T}, nodes_to_remove). Since remove_nodes(g, p) calls the same method, node dropping on a batch of graphs hits this too.

I think the fix is just to index the indicator with the kept nodes (g.graph_indicator[nodes_to_keep]) like is already done for ndata.

GNNGraphs master (17d1906, 1.6.0-DEV), Julia 1.12.1