Callbacks with parameters.

[erick-xanadu]

Context: Callbacks with parameters. Second step in implementing callbacks. Next step will be return values.

Description of the Change:

1. Collect type information and tree shape during tracing
2. Adds parameters to PythonCallOp.
3. Bufferizes said parameters.
4. Lowers PythonCallOp to an LLVM::CallOp where the first parameter is the function pointer, the second parameter is the number of parameters, and the rest of the parameters are pointers to memrefs.
5. In the runtime, use va_list to collect all memrefs.
6. Pass va_list to pybind module
7. In pybind moulde iterate over va_list and add these to a python list
8. In python, use the previously collected metadata to cast va_list pointers to jax arrays.
9. Unflatten and pass the result to the user's callback.

Benefits: Callbacks with parameters.

[sc-58870]

[rmoyard]

It looks good to me! Mostly questions for my understanding.

Do we agree that there is no restriction on the arguments pytree structure? It could be dictionnaries for example?

[erick-xanadu]

@rmoyard I don't think build wheels are required here. I am not changing the build system. The wheels were checked in the previous merge.

[erick-xanadu]

Do we agree that there is no restriction on the arguments pytree structure? It could be dictionnaries for example?

As long as they implement the pytree interface, they should work. The tests here include a dictionary.

[rmoyard]

great job 🥇