jax

7d924e8f - Make canonicalization of reshapes more robust:

Make canonicalization of reshapes more robust: - do a little graph traversal to find the load op - handle more leading dims - single dim reshape fallbacks PiperOrigin-RevId: 774471363