jax

a2d67a6a - create a mosaic_gpu_support cc_library and include it as part of the //third_party/py/jax/jaxlib/cuda:cuda_gpu_kernels