Using Scatter cpp implementation doesn't enable gradient tracking ?

[Kheims]

Hello everyone,

First of all, I wish all of you a great year filled with joy and success.

To put some context, I am a Phd Candidate working on Graph Neural Networks, and was really impressed by MLX.

I was trying the past few days to implement Graph Neural Networks MessagePassing's paradigm in MLX (following pyg's official implementationand other renown papers advocating the efficiency of using operations such as scatter and gather in this particular context).

To do it for Graph Convolutional Network for example the use of scatter_add implementation is required, therefore I have binded that cpp function implementation to python, I was able to use it in python and it has passed the equivalent cpp tests, the forward pass has a correct behaviour but during backward the grads aren't updated.

From my understanding, when using an operation such as scatter_add (that in short term just uses scatter with a Reduce Type of Sum) the primitive needs to implement jvp and vjp, as the operation itself builds the computation graph but doesn't update parameters.

The scatter function in the file mlx/primitives.cpp doesn't have implementation of jvp and vjp and therefore can't be used yet, am I correct ?

TLDR :

• Is the scatter function with it's different strategies (scatter_add, _max ...etc) usable as of now in python API if we just provide a binding to it ?

• If we want to make core contributions to operations and primitives ( say for example implementations of jvp and vjp ) is there a more efficient way to isolate the functions ( advices from your personal experience would be great :D ) rather than needing to modify the code and rebuilding all the project.

• I have read the developer documentation, but if you have any other resources that might be useful that would be amazing.

[awni]

The scatter function in the file mlx/primitives.cpp doesn't have implementation of jvp and vjp and therefore can't be used yet, am I correct ?

Correct (open issue for this is #282).

Is the scatter function with it's different strategies (scatter_add, _max ...etc) usable as of now in python API if we just provide a binding to it ?

Indexing uses scatter already. We could bind scatter_* if there is need and then it would be usable in Python.

If we want to make core contributions to operations and primitives ( say for example implementations of jvp and vjp ) is there a more efficient way to isolate the functions ( advices from your personal experience would be great :D ) rather than needing to modify the code and rebuilding all the project.

If I understand you question correctly: cmake should mostly handle not rebuilding the whole project with every change. It only rebuilds the .o for the .cpp you update.

I have read the developer documentation, but if you have any other resources that might be useful that would be amazing.

Not at the moment, but if you keep track of your pain points and we can update accordingly!

[Kheims]

Thank you for your answers.
How do we usually test our implementations of jvp and vjp ? The naive way I would do it, is rebuilding project and then in mlx side use mx.grad to make sure it coincides with numerical results. Is there a way to reduce this overhead and iterate more quickly ?
I saw unit tests for operations in ops_tests but I don't have an idea of how to test jvp and vjp in cpp side, are there exemples in the code base ?

[awni]

The naive way I would do it, is rebuilding project and then in mlx side use mx.grad to make sure it coincides with numerical results. Is there a way to reduce this overhead and iterate more quickly ?

Every change you make in C++ requires a rebuild 😄 . But it shouldn't be that slow. Usually I can rebuild in seconds for a small change to a .cpp file.

I saw unit tests for operations in ops_tests but I don't have an idea of how to test jvp and vjp in cpp side, are there exemples in the code base ?

Exactly the same as in the python side but using the C++ functions instead. You can look in autograd_tests.cpp

[TristanBilot]

Hi @awni and @Kheims, I jump into the conversation as I'm working with @Kheims on the implementation of a message passing neural network architecture to include within MLX.

At the moment, we want to implement vjp and jvp for the scatter_add function.
After modifications within the cpp code, we use to recompile the project with env CMAKE_BUILD_PARALLEL_LEVEL="" pip install . but we think that this is not efficient and we're asking which commands do you usually use to recompile your modified files?

Also, do you use any debugger during your developments?
Thanks! 🙏

[awni]

This is just me, others may have other solutions:

• Usually when I develop the C++ part of the code I start with C++ tests and a C++ only build. Instructions on how to build the C++ only library

• You can run the C++ tests from the build/ directory and specify the test you want like this ./tests/tests -tc="my test name"

• So my flow with C++ is like make -j && ./tests/tests -tc="my test name" to see if my change fixes a test

• gdb + std::cout for debugging, I don't have a great answer there. From C++ it's a bit tricky as you need an eval to see the result of an array

• If you're doing python dev you should use an editable install: env CMAKE_BUILD_PARALLEL_LEVEL="" pip install . -e.

[TristanBilot]

Thanks a lot, this is exactly what we were looking for!

[awni]

As a side comment, are you certain you need the vjp of scatter_add? Maybe there is a way to do what you want with the existing ops? If you post a snippet here we can take a look.

[awni]

By adjacency matrix I mean $A \in \mathbf{R}^{|V| x |V|}$ with a 0 or 1 in $A_{ij}$ if there is an edge between $i$ and $j$. Then you multiply $A$ by the messages.

You can construct $A$ in numpy / python in your data processing pipeline. My hunch is this would also be more efficient up to pretty large graphs (maybe a few hundred nodes, maybe even more) compared to a scatter add. But I'm not certain there. But certainly an easier starting point and you can then go back and update it with scatter_add if it's a bottleneck.

[Kheims]

Our standard benchmark will be on Planetoid Dataset specifically on Cora Dataset that contains 2708 nodes and approximatively 13k edges (including self edges).
Our prior aim was to provide an interface to use GNNs similar to pyg's one, providing GCN, GAT and other modules built on top of the MessagePassing interface (in short, our initial intent was purely contribution) . I agree on your comment that for small graph it might be better to directly use an adjacency matrix, perhaps we can use strategy here to internally use the right operation depending on the provided graph (ie #num_of_nodes).
The reason for using scatter is induced by the fact that we want to deal with worst case scenarios where there is a large graph and connectivity is sparse, @TristanBilot has already benchmarked GCN in MLX without using scatter, but we wanted to inspect this path too.

[TristanBilot]

Using a dot product with a dense adjacency matrix is indeed what we've done in the GCN example. You are right when you say that it is fast and sufficient for small graphs.

However, only few GNNs can be implemented using a simple adjacency matrix. In practice, GNN frameworks use the message passing paradigm, which allows to build almost any GNN based on the scatter operation.
We want here to propose within mlx an abstract class MessagePassing that anyone can implement to build any GNNs on top of it, similarly as in torch_geometric (pytorch) and jraph (jax).

We think that providing this new feature within mlx would be very useful for graph ML researchers doing experiments on their Mac without requiring expensive GPUs.

[awni]

Great! I fully support using scatter add and implementing the grads you need! Just wanted to make the suggestion in case you hadn't considered it.

The other thing I notice is the following which one has to be careful of as it will eval the graph. I assume the indices are inputs to the mode (not computed on the fly). In that case it should be ok.

def _unique(array):
    return len(set(array.tolist()))

[TristanBilot]

Ok we'll be careful with this function thanks ;)
We stay in touch with you when we progress in the implementation!

[awni]

I am not sure to have yet understood why we need both jvp and vjp for the backpropagation, maybe can you help us understand if we need both in our scenario?

For backprop you just need VJP. Usually we implement both because they end up being pretty similar but in your case you can just do VJP. Throw std::runtime_error("[scatter::jvp] Not yet implemented"); for JVP if you don't implement it.

And shall we create specific cases in jvp and vjp for each reduce_type (add, mean, etc)?

No you don't need to do all of them, you can just do the one you need and throw for the ones you don't. But feel free to add them if you're interested / up for it!

[angeloskath]

Sorry to be late to the party but there is a branch (scatter-vjp I think it is called) where a few weeks ago I started implementing the vjps of scatter.

Simple scatter is already there working and scatter add shouldn't be too hard to implement in a similar fashion.

[TristanBilot]

Hi @angeloskath, I'm back in this conv as I'm looking to implement the vjp for scatter_max.
My first intuition is:

• grad for source is cotan at indices where the source is the max value
• grad for inputs is gathering cotan for indices where the inputs are the max value

Can you please correct me if I'm wrong?
Another question is how to keep track of the indices which are the max values during the forward?

[angeloskath]

Your first intuition is exactly right. So this would be in mlx pseudocode:

result = scatter_max(...)
src_vjp = cotan * (src == result)
upd_vjp = gather(cotan) * (gather(result) == upd)

[angeloskath]

completely random but you can also check all these gradients using the assignment gradient if there are no double indices and all the updates are unique.

Meaning that

x[idx] = mx.maximum(x[idx], updates)

and

x = x.at[idx].maximum(updates)

should have the same gradient under the above assumptions.

[TristanBilot]

How can I access result from the vjp? We only have values, updates and indices from primals.

[angeloskath]

You can just "recompute" it. Simplify will take care of reusing the result.

[gboduljak]

Hi @Kheims, @TristanBilot :) @francescofarina recently opened #380, aiming to introduce GNN layer. We also recently started working on Python bindings for scatter_{ops} (#391). I think it is a good idea you open a draft PR for your message passing layer so we do not duplicate the work.

[TristanBilot]

Hi @gboduljak, nice to meet you!
Thanks for telling us about this PR, not sure we'd see it without you :)
As suggested, I created this PR #398 for the implementation of the basic MPNN architecture. We can continue the discussion about the integration of GNNs within this new PR.