JAX vs TensorRT

Autonomous Machines Jetson & Embedded Systems Jetson TX2
1

Hi

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?

3

Hi,

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.

Thanks.

1 Like