Helper functions for Graph permutation

[InnocentBug]

One of the fundamental concepts of graph based data is, that changing the order of the nodes (or the edges) does not change the graph properties.
Basically all GNNs work can and maybe should be tested if the results are invariant (or at least equivariant) under such permutations.

And I thought it might be helpful if we had a way to apply such a permutation quickly to the graphs and see if that is true.

I needed some helper functions for my own work that applies such permutation, and since I appreciate this open-source project, I was wondering if you would like to include them in the official jraph repo.

I implemented 2 helper functions that can apply permutations in the nodes, and permutations in the edges.
One fundamental requirement I set for myself was, that this should still work for batched and padded graphs, and the permutations should be easily applied only inside the individual graphs inside a graph tuple.

I added a user friendly way to achieve this randomly, but also have an option for more controlled permutation.
As for example it is needed in order to invert a permutation.

I am open for feedback on this design decision.

JIT

Unfortunately, I can't quite figure out a way to write the node-wise permutation sequence generation in a way that is jitable.
I don't this it fundamentally necessary, but it would nice if we could make that possible. I am open for ideas.
The problem, I am running into is, that jax.random.permutation cannot be vmap for variable sizes of the permutation.

Tests

I added some test for the helper functions.
At the moment it is more an end-to-end test then a unit test, so if you have suggestions I am willing to improve there.

Also, I am using the helper function to generate random graphs to test with.
I needed a bit more control about the size of individual graphs.
With too small graphs, the probability of having a permutation that doesn't actually do anything was too big, so the test would falsely detect that the permutation didn't work.

Generally, I since my test is build on random graphs with random permutations, a little brittleness remains.
Let me know if I should switch to fixed graphs instead to eliminate the brittleness.

I hope this a contribution you would like to include.

[InnocentBug]

Also, in the future, I was thinking about including some code for measures on graphs.
And this was the first step, since a metric is required to be invariant under permutation.

Let me know if such a contribution would be welcomed as well.

[InnocentBug]

@jg8610 is there interest in this PR?