Would you consider including JAX into LambdaStack?
We are looking into adding Jax. It requires more testing.