jax

218560df - #sdy fix-forward jaxlib compatibility issue

#sdy fix-forward jaxlib compatibility issue With upcoming shardy flip in jax the incompatibility issue was uncovered with jaxlib extension for xla-shardy. PiperOrigin-RevId: 784636442