jax

511bf2fe - Move jax._src.lax to its own BUILD rule

Move jax._src.lax to its own BUILD rule This also bundles-in `ad_checkpoint.py` and `state/*.py` because these have circular dependencies on lax source files. Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, prevents use of internal APIs, and leads to improved build and iteration times. PiperOrigin-RevId: 773030035