jax

3f91b4b4 - Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/

Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/ Move the common jaxlib/gpu_plugin_extension into jaxlib/gpu/ Cleanup only, no functional changes intended. PiperOrigin-RevId: 738183402