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