jax

74e912c3 - #sdy dynamically choose which `custom_partitioning` API to use based on the current

#sdy dynamically choose which `custom_partitioning` API to use based on the current value of the `use_shardy_partitioner` feature flag. Before the way the API works depends on the value of the flag when the partitioning is defined. But we should allow this to be dynamically swapped in and out when the function is actually called. This change allows for that. PiperOrigin-RevId: 715293018