jax

d66be780 - [shape_poly] Add support for shape polymorphism with native serialization for lax.linalg.qr on TPU