Its really cool! If you don't mind me asking, does it support variable size inputs? I am bit confused about JAX in that regards. I am trying for long to run JAX stablehlo models in C++ for inference. However dynamic shapes were still an issue. If I understand correctly, it recompiles the kernels for all different shapes at runtime, so if inputs vary too much in shape, it will spend considerable time recompiling kerenels. In C++ inference it becomes impossible .However I could be wrong (I did not fully understand the issue, the developer of gopjrt tried to explain it to me!). Do you have any thoughts on this?
e.g.:
https://github.com/openxla/xla/issues/33092 https://github.com/openxla/xla/issues/35556
Explanation from the gopjrt dev:
https://github.com/gomlx/gopjrt/issues/59