JAX vs TensorRT


I would like to know if someone has testes JAX in an embedded GPU?
JAX (JAX reference documentation — JAX documentation).

Jax accelerates training and inference in cloud GPUs and TPUs. I do not know if it is faster than TensorRT or if it works with embedded GPUs and TPUs.

Has anyone tested it?


Sorry that we don’t have too much experience on JAX.

But it seems that JAX is a language-level optimizer.
TensorRT is an algorithm-level optimizer and should give you a better performance.


1 Like