jax

fcaced32 - Copybara import of the project:

Copybara import of the project: -- ced333d1d4aec2825e9afd81c2ca9721b7e3cc67 by Matthew Johnson <mattjj@google.com>: redo #4535 lazy simplification PiperOrigin-RevId: 338606348