jax

ba8c3a92 - Fix missing batch partitioning hook in GPU plugin.

Fix missing batch partitioning hook in GPU plugin. I had somehow missed properly registering the "batch partitionable" registration hook on the GPU plugin, causing a segfault when the missing pointer was accessed. This fixes that and updates the tests to make sure that the registration code is executed even without multiple devices. PiperOrigin-RevId: 725312473