[JAX] Update users of deprecated jax.lax.* submodule exports.
[JAX] Update users of deprecated jax.lax.* submodule exports.

jax.lax.lax/lax_parallel/lax_control_flow were internal names that will no longer be exported. Update users to use either jax.lax.* public names, or to explicitly use the new internal names (jax._src.lax.*).

PiperOrigin-RevId: 337913333

tensorflow/probability,tensorflow/probability