JAX for inference only: what are you seeing?

[edgecoast]

I think the bottleneck is somewhere else. I wouldn't optimize around that first. On my side the bigger swing came from the runtime and not the model. Could be a workload difference, but I wouldn't generalize too fast from one setup.