Processing Block - JAX

Hello EI-Team,

@janjongboom mentioned at EI Image 2022 JAX. (If I read the Quickstart it is Numpy for CPU, GPU and TPU.) @janjongboom said it is a Python to graph compiler and gives the possibility to translate Python code to Tensorflow Lite and finally to source code, by using the EON tuner. Correct?
Am I also correct that the final goal is to get a Tensorflow Lite signal processing block that we can plug into together with your neural network block into the code?

I am starting to write a custom processing block to extract specific features, together with an event detection algorithm. I am using mainly numpy and scipy. Does it make sense to use JAX where it is possible? If yes, what do I need to take into account during the coding?

Thanks for the feedback.

Regards,
Joeri

Hello @Joeri,

I have not tested writing a custom DSP block using JAX so I am not yet 100% confident.
However, I think you have the correct approach.

Does it make sense to use JAX where it is possible?

I would say yes, definitely.

what do I need to take into account during the coding?

I will let someone else more knowledgable fully answer your questions

Hi @Joeri , yeah this is the plan - but we haven’t finished it yet :slight_smile: A good way to test would be to follow example-custom-ml-block-scikit/train.py at 0f7582b5e6bd535d6142e85487a73c68df7b217a · edgeimpulse/example-custom-ml-block-scikit · GitHub - just throw your DSP function through the conversion function and see if it succeeds.

We’re currently building some building blocks in the SDK to make it easier to port these functions, and then hook it into EON Compiler to get on-device source code.

Thanks for the feedback.

I am also not finished yet with the custom DSP block :slightly_smiling_face: But I am early phase of development, so it is a good idea to make the blocks future-proof.
Keep me posted!