Jax and real-time data acquisition

[Filipe-Camarneiro]

Hi, first time here.

Is it possible to use Jax for optimization routines using real-time data acquisition?

I have implemented optimization routines based on propagation methods with Jax and they work properly. Now i want to substitute the propagation methods with real-time data acquisition.

The idea is that starts with a MxN matrix as a guess - this matrix is used to acquire data - and the loss function is calculated based on the data acquired. The MxN is then optimized by Jax which is then utilized again for data acquisition and so on (more or less a hardware on the loop concept).

Thank you for your time.

[jakevdp]

I've never seen anything like this, but you could probably use one of two strategies:

• do your optimization loop in Python: this would allow you to acquire data via Python between optimization steps
• do your optimization loop in JAX, but use io_callback between steps to acquire data.

Performance-wise, both options are probably comparable. If you use the first, you should ensure that your data shapes stay constant from iteration to iteration in order to avoid recompilation at each step. In the second, that condition is required by the nature of io_callback.