jax

18259e7c - Add modes for `config.jax_dump_ir_to` via comma-delimited `config.jax_dump_ir_modes`.

Add modes for `config.jax_dump_ir_to` via comma-delimited `config.jax_dump_ir_modes`. - If `jax_dump_ir_to` is specified, dump mlir by default. - Also gives options to dump JAXPR via `jax_dump_ir_modes=jaxpr` and an equation count profile for pprof via `jax_dump_ir_modes=eqn_count_pprof`. PiperOrigin-RevId: 776136562