jax

f380c3c3 - Merge pull request #20179 from keshavb96:custom_gpu_ops_partitioning