Question Call Python (JAX) code from unity?

Hi,
I want to develop a game that heavily relies on GPU simulation. I am prototyping it in Python (the simulation part) with machine learning frameworks (since I know them quite well). In particular JAX.

I know there are versions of Tensorflow and Pytorch for C#. After many failed attempts I even managed to pull in one of them (Pytorch) into Unity without making the whole thing crash. But still, I am way less productive than before.

Does anyone know if there is a way to share the memory between the Python script and the Unity script so that I basically could call ```step_simulation``` once every few frame, let the stuff be done by Python and then just use the variables from C#? Another problem would occur since the JAX arrays will all be on the GPU, so I'd need to handle that somehow.

Any pointers into what I should look for would be great!
I know that in the Godot engine you are able to do similar things (allthough I am not sure about the GPU arrays). But I prefer Unity and I am a huge fan of the DOTS ECS, which I think would complement the GPU simulation well by letting me also fully utilize the CPU in a parallel fashion.

 

**Bump**