jax

d73557de - add offload-compress flag to reduce rocm wheel sizze